Update model.py

musecoco/1-text2attribute_model/model.py

@@ -436,7 +436,7 @@ class BertForAttributModel(BertPreTrainedModel):
             i += 1
             logits[k] = self.classifieratt[idx](self.dropout(pooled_outputs[k]))
 
-        total_loss = torch.Tensor(0).to(input_ids.device)
+        total_loss = torch.Tensor(0.0).to(input_ids.device)
         loss = {}
         if labels is not None:
             for k in labels.keys():