Could you check the shapes of wlast and fun? You might be broadcasting the tensors (which might not be what you want) and thus the memory could increase significantly.
Here is a small example:
x = torch.randn(10, 1)
y = torch.randn(1, 10)
diff = x - y
print(diff.shape)
> torch.Size([10, 10])