Hi, I’m using nn.DataParallel(model). The training is working well. But After some epochs; i do some computation which seems to be happened on one GPU and thus a CUDA out of memory is thrown. This is the part of the code that gives the error
  File "/home/xxx/xx/xxx/trainers.py", line 115, in compute_distance
    dists.add_(torch.tril(100000 * torch.ones(len(features), len(features)).cuda()))
RuntimeError: CUDA out of memory. Tried to allocate 3.96 GiB (GPU 0; 11.17 GiB total capacity; 9.46 GiB already allocated; 1.39 GiB free; 17.78 MiB cached)

Is there a way to make add_ op works. I don’t want to reduce the batch-size as this will increase the training time.
Thank you