Unset CUBLAS_WORKSPACE_CONFIG for non-deterministic training and inference (#19138)

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>

ultralytics/engine/model.py

@@ -140,6 +140,7 @@ class Model(torch.nn.Module):
             return
 
         # Load or create new YOLO model
+        __import__("os").environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"  # to avoid deterministic warnings
         if Path(model).suffix in {".yaml", ".yml"}:
             self._new(model, task=task, verbose=verbose)
         else:

ultralytics/engine/trainer.py

@@ -52,6 +52,7 @@ from ultralytics.utils.torch_utils import (
     select_device,
     strip_optimizer,
     torch_distributed_zero_first,
+    unset_deterministic,
 )
 
 
@@ -471,6 +472,7 @@ class BaseTrainer:
                 self.plot_metrics()
             self.run_callbacks("on_train_end")
         self._clear_memory()
+        unset_deterministic()
         self.run_callbacks("teardown")
 
     def auto_batch(self, max_num_obj=0):