I can't compute loss properly

[Rita_Pucci]

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?