I have worked for pytorch with official released version but there were some limits related to GPU. For example, executing torch.range() on GPU gave me the below error,

torch.cuda.FloatTensor doesn't implement stateless method arange

Therefore I just cloned pytorch git and build the docker image from the provided dockerfile.
Finally, the code execution of torch.range() on GPU was not the problem but I got an another issue.

RuntimeError: cuda runtime error (2) : out of memory at /tmp/pip-ikrnhryb-build/torch/lib/THC/generic/THCStorage.cu:66

Firstly, I tried to change batch_size from 16 to 8 but gave me the same error.
I tried batch_size 2 and 1 as well but result in with the below errors respectively.

invalid argument 2: sizes do not match at /tmp/pip-ikrnhryb-build/torch/lib/THC/generated/../generic/TH --> batch_size = 2

assert len(modules) == len(devices) --> batch_size = 1

Since now there is a memory issue with the latest version, I am wondering if it will take some times to be completely stable. Or is there any way to resolve this situation?

Need your help!