I want to write a loss function which will create new variables according to the shape of input data. The matter is when I use my loss with loss.cuda(0), the data in it won’t move to gpu device, and if the input data is on gpu it will occur error.
I referred to the code of loss functions already exists in PyTorch and noticed there is a class named _WeightedLoss and is used in NLLLoss. But because the weight NLL use is created in its init and needn’t adjust according to the input data, I don’t think this way is suitable for my problem.
The function I want to write can approximately described like below:
class exampleLoss(nn.Module):
def __init__(self):
super(exampleLoss,self).__init__()
def forward(self,input, groundTruth):
data1=torch.zeros_like(input)
#some computation
return result
I also noticed that the net I wrote before which also inherit from nn.Module can use net.cuda() to put all the layers and parameters to gpu even I don’t use functions like nn.ModuleList or nn.register_buffer. What’s the mechanism PyTorch used for collecting members of one class and make them available for cuda() operation?
Briefly, I want to use .cuda() to put my loss function which will create variables during forward() to gpu and is curious about the mechanism behind Module.cuda() operation.
Any help is appreciated. Thanks!
