a strange problem about saving memory in training process

[zhang-wen]

hello guys,

i have some strange problems here,

if i use memoryEfficientLoss function to backward loss, the training seems to be normal

but if i put the content of memoryEfficientLoss into trainEpoch function and do not define extra memoryEfficientLoss function, the training will not converge, all other code are the same.

and another question is that i guess the split operation along the first dimension to ouputs of the model can save the memory, however, if so, how do we calculate gradients and do backward and why do you call the backward() two times (loss.backward() and outputs.backward()) ? can you explain this ? thank you.

Can anyone tell me why ? any reply will be appreciated.

[srush]

Check out the new version of the code. I think it makes this function much clearer.