Hi, Thanks for the reply and the function reference. I used it and found something confusing. I have two RTX 2080 Ti gpus. The every third number shows the cache size at that point using the function torch.cuda.max_memory_cached ( )
When I use gpu 0, get the following output,
0 6.981476306915283
torch.Size([50, 1000])
10034.0
1 6.936072826385498
torch.Size([50, 1000])
10034.0
2 7.297829627990723
torch.Size([50, 1000])
10034.0
3 7.23666524887085
torch.Size([50, 1000])
Traceback (most recent call last):
  File "test.py", line 44, in <module>
    loss_.backward()
  File "/home/ahosain/workspace/dl_env/lib/python3.6/site-packages/torch/tensor.py", line 166, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/home/ahosain/workspace/dl_env/lib/python3.6/site-packages/torch/autograd/__init__.py", line 99, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: CUDA out of memory. Tried to allocate 1.50 GiB (GPU 0; 10.76 GiB total capacity; 5.23 GiB already allocated; 1.48 GiB free; 3.19 GiB cached)

In case of gpu 1,  it seems running fine with following output,
0 7.1220622062683105
torch.Size([50, 1000])
8964.0
1 7.066858291625977
torch.Size([50, 1000])
9512.0
2 7.112499237060547
torch.Size([50, 1000])
9512.0
3 7.206641674041748
torch.Size([50, 1000])
9536.0
4 7.3534345626831055
torch.Size([50, 1000])
9536.0
5 7.341899871826172
torch.Size([50, 1000])
9536.0
6 7.346314907073975
torch.Size([50, 1000])
9536.0
7 7.240469932556152
torch.Size([50, 1000])
9536.0
.
.
.
.