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.register_buffer. What’s the mechanism PyTorch used for collecting members of one class and make them available for
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
Any help is appreciated. Thanks!