Fix loss for a classifier

pycroscopy · Apr 13, 2023 · 00f2262 · 00f2262
1 parent f8266bc
commit 00f2262
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/atomai/trainers/trainer.py b/atomai/trainers/trainer.py
@@ -971,6 +971,8 @@ def __init__(self,
         kwargs["batch_seed"] = kwargs.get("batch_seed", seed)
         set_train_rng(seed)
         self.nb_classes = nb_classes
+        self.criterion = self.get_loss_fn('nll')
+
         self.net, self.meta_state_dict = init_cls_model(
                                 nb_classes, backbone, **kwargs)
         self.net.to(self.device)