Custom training

I want to make my training process with simpe CNN (on MNIST): first iteration - forward on training set, backward on 80% of training set where loss value was highest on first, at the second iteration, a forward pass is made for those 80% from the previous backward pass, 80% of examples with the highest loss value are taken from them, these 80% of 80% are used for the backward pass; etc. At every fifth iteration, we start with the entire dataset.
Any ideas how to implement this?
Thank you in advance.

You could use an unreduced loss (via passing reduction='none' to the criterion) and sort the loss values (and store the sort indices). Once they are sorted you could slice them, reduce the loss, call backward` on it, and update the parameters.
In the next step you could reuse the sort indices from the previous run and repeat the procedure using your schedule.