Out of memory issue while running ''def _test_single_epoch(self):'' function

[Gao Le Dai]

Hi there, I'm currently trying to reproduce the training and evaluation process but have run into an issue with the "def _test_single_epoch(self):" function causing an OOM error.

I have this issue running on a 24G RTX4090 GPU, it seems that this is not due to hardware limitations.

Do you have any suggestions on how to modify the function and release some memory during the evaluation?

########################################################
Trying to train from scratch

########################################################
modify

get

confirm the issue is due to ''def _test_single_epoch(self):'' function
########################################################
solely evaluating with pre-trained weight