Description
Hi, @soumyadeepdey
When I use your code to train a model in gpu, it seems always occer the problem of OOM.
The batchsize is 1, other parameters haven't changed.
Please check it,thanks
In which platform does it happen?
linux
gpu
How do we replicate the issue?
You can use the command python3 sample_train.py, you can replicate the issue.
Expected behavior (i.e. solution)
Other Comment

s
