fixed model training

2025-11-12 18:11:19 +02:00
parent 7cb4201bc0
commit 8354aec830
6 changed files with 9 additions and 9 deletions
--- a/NN/models/cob_rl_model.py
+++ b/NN/models/cob_rl_model.py
@@ -323,7 +323,7 @@ class COBRLModelInterface(ModelInterface):
        self.optimizer.zero_grad()
        
        if self.scaler:
-            with torch.cuda.amp.autocast():
+            with torch.amp.autocast('cuda'):
                outputs = self.model(features)
                loss = self._calculate_loss(outputs, targets)