CudnnBatchNormBackward nan temporary fix ultralytics#12

:)
manole-alexandru · Apr 10, 2023 · ef08e87 · ef08e87
1 parent 2975086
commit ef08e87
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/train.py b/train.py
@@ -338,7 +338,10 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
                     loss *= 4.
 
             # Backward
-            scaler.scale(loss).backward(retain_graph=True)
+            try:
+                scaler.scale(loss).backward(retain_graph=True)
+            except Exception as e:
+                print('\n-------- FOUND EXCEPTION: ', e, ' Life goes on.-------\n')
 
             scaler_seg.scale(loss_seg).backward()