Thanks but I already solved it. I have another question. Why are my tensors have extra brackets like below one? Am i doing something wrong or it is ok? I am new in pytorch and still learning as you can see. Thanks.
In [35]: distfunc
Out[35]:
tensor([[[[0.7022, 0.6965, 0.6908,  …, 0.6956, 0.7012, 0.7069],
[0.6995, 0.6937, 0.6879,  …, 0.6927, 0.6984, 0.7040],
[0.6967, 0.6909, 0.6851,  …, 0.6899, 0.6955, 0.7012],
…,
[0.6967, 0.6909, 0.6851,  …, 0.6899, 0.6955, 0.7012],
[0.6995, 0.6937, 0.6879,  …, 0.6927, 0.6984, 0.7040],
[0.7022, 0.6965, 0.6908,  …, 0.6956, 0.7012, 0.7069]]]],
dtype=torch.float64)