Would you like to add DropoutFC somewhere into Train?
Maybe I’ve missed it, but it seems not being used.
Is your model overfitting and you would like to regularize it with your module?
As a side note, it looks like your Train class might store the computation graph in every iteration thus leading to an increasing memory which finally might yiel an out of memory error.
If you need to store the loss values for debugging, you should just store the python float value instead of the tensor. Use this line instead:
total_loss += loss.item()

PS: I’ve formatted your post, since the code was not readable. You can add code blocks using three backticks.