Thanks Spandan. The thing that I was trying:  there is a main branch of network which takes input images. At some point inside the network, I would like to learn another probability (a tensor with a distribution and a linear upsamling + convtranspose).
When I create Tensor inside the network, it cannot decide cuda/cpu. Thus, I give it as a second input and not automatically change all parts inside the net definition to GPU/CPU and torch.mm 37 or torch.mul won’t be an issue.