fix args.loss.obj AttributeError

mindspore-lab · Oct 17, 2024 · 70520af · 70520af
1 parent eeebe04
commit 70520af
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -189,8 +189,9 @@ def train(args):
         eval_dataset, eval_dataloader = None, None
 
     # Scale loss hyps
-    args.loss.cls *= args.data.nc / 80
-    args.loss.obj *= (args.img_size / 640) ** 2
+    args.loss.cls *= args.data.nc / 80  # scale to classes and layers
+    if hasattr(args.loss, "obj"):
+        args.loss.obj *= (args.img_size / 640) ** 2  # scale to image size and layers
 
     # Create Loss
     loss_fn = create_loss(