I can't compute loss properly

[Riretta]

Good morning everyone, I have a model defined as: def init(self,NUM_CLASSES): super(model, self).init()

Features function :

    self.modelCaps1 = ResNetCaps.ResNetCaps(NUM_CLASSES)
    self.modelCaps2 = ResNetCaps.ResNetCaps(NUM_CLASSES)
    #Classification function:
    self.modelLin = nn.Linear(NUM_CLASSES, NUM_CLASSES)

def forward(self,inputs):
    digit1,_ = self.modelCaps1(inputs)
    digit2,_ = self.modelCaps2(inputs)
    output = self.modelLin(self.bilinear(digit1, digit2))

I trying to figure out how to compute the loss in this case. My first attempt was a simple loss: criterion = nn.CrossEntropyLoss() loss = criterion(x,target) loss.backward()

but after few epochs i got nan

Then I'm trying with a customized loss: loss_SM = criterion(x,target) A = digit1 B = digit2 def model_loss(self, loss_SM): loss_caps1 = self.modelCaps1.model_loss(self.A,torch.sum((self.A[:,:,:,:]loss_SM),dim=2).squeeze()) print(loss_caps1) loss_caps2 = self.modelCaps2.model_loss(self.B,torch.sum((self.B[:,:,:,:]loss_SM),dim=2).squeeze()) print(loss_caps2) loss = loss_SM+ loss_caps1 + loss_caps2 but in this case i have -316

Where i'm doing wrong?