InnovArul:

can you change this portion of the code to the following and try once?


I tried it but that did not solve the problem. (I wonder why.) However I found a promising solution: I changed the dataset to return tensors that do not require gradients and change require_grads after getting a batch from the dataloader. I have not yet integrated everything so it might be too early for calling it a solution.