second

scripts/training/trainer.py

@@ -44,8 +44,11 @@ class NarrativeTrainer:
         self.global_step = 0
         self.best_val_f1 = 0.0
         
-        # Initialize mixed precision training
-        self.scaler = GradScaler('cuda', enabled=self.config.fp16)
+        # Initialize mixed precision training (Fixed version)
+        if self.config.fp16:
+            self.scaler = torch.cuda.amp.GradScaler()
+        else:
+            self.scaler = None
         
         # Setup training components
         self.setup_training()
@@ -204,7 +207,8 @@ class NarrativeTrainer:
         for step, batch in pbar:
             batch = {k: v.to(self.device, non_blocking=True) for k, v in batch.items()}
             
-            with autocast(enabled=self.config.fp16):
+            # Mixed precision training
+            with torch.cuda.amp.autocast(enabled=self.config.fp16):
                 outputs = self.model(
                     input_ids=batch['input_ids'],
                     attention_mask=batch['attention_mask'],
@@ -213,17 +217,27 @@ class NarrativeTrainer:
                 loss = self.criterion(outputs, batch['labels'])
                 loss = loss / self.config.gradient_accumulation_steps
             
-            self.scaler.scale(loss).backward()
+            # Backward pass with scaler if fp16 is enabled
+            if self.config.fp16:
+                self.scaler.scale(loss).backward()
+            else:
+                loss.backward()
             
             if (step + 1) % self.config.gradient_accumulation_steps == 0:
-                self.scaler.unscale_(self.optimizer)
+                if self.config.fp16:
+                    self.scaler.unscale_(self.optimizer)
+                    
                 torch.nn.utils.clip_grad_norm_(
                     self.model.parameters(),
                     self.config.max_grad_norm
                 )
                 
-                self.scaler.step(self.optimizer)
-                self.scaler.update()
+                if self.config.fp16:
+                    self.scaler.step(self.optimizer)
+                    self.scaler.update()
+                else:
+                    self.optimizer.step()
+                    
                 self.scheduler.step()
                 self.optimizer.zero_grad()