I need to use elementwise mutliplication (torch.mul) and matrix multiplication (torch.mm) inside model definition in forward pass in CPU/Cuda. Is there any CUDA version of these operations which can be used w/o switching into CPU mode?
The operations don’t reside on cuda, the tensors do. When you move your model to CUDA using model.cuda() and pass into the model function tensors say input and target that are on cuda() all computations will take place on GPU I believe.
So you can define your model using just torch.mm or torch.mul, just make sure that the tensors/variables that are passed into the model() are on cuda and that your model is also on cuda.
Hope this helps!
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 or torch.mul won’t be an issue.
If you create a new tensor inside the model function, it will be on cuda if the model is on cuda.