Solved by spanev in post #2 
              
              
                Hi @jpainam, 
You have to move the torch.ones Tensor to the GPU: 
dists.add_(torch.tril(100 * torch.ones(len(features), len(features)).cuda()))