Fix optim import and ensure PyTorch AMP GradScaler backwards compatibility

train.py

@@ -21,6 +21,8 @@ Key Features:
   * Tokenizer files saved alongside ONNX models for standalone portability.
 """
 
+from __future__ import annotations
+
 import os
 import sys
 import re
@@ -30,14 +32,15 @@ import shutil
 import tempfile
 import urllib.request
 import argparse
-from typing import Tuple, Dict, Any
+from typing import Tuple, Dict, Any, List
 
 import numpy as np
 import pandas as pd
 import torch
 import torch.nn as nn
-from torch.utils.data import Dataset, DataLoader
+import torch.optim as optim
 from torch.optim import AdamW
+from torch.utils.data import Dataset, DataLoader
 from transformers import (
     AutoTokenizer,
     AutoModelForSequenceClassification,
@@ -151,7 +154,7 @@ def train_epoch(
     dataloader: DataLoader,
     optimizer: optim.Optimizer,
     scheduler: Any,
-    scaler: torch.amp.GradScaler,
+    scaler: Any,
     criterion: nn.Module,
     device: torch.device,
 ) -> float:
@@ -167,7 +170,12 @@ def train_epoch(
         labels = batch["labels"].to(device)
 
         if use_amp:
-            with torch.amp.autocast(device_type="cuda", dtype=torch.float16):
+            try:
+                autocast_ctx = torch.amp.autocast(device_type="cuda", dtype=torch.float16)
+            except (AttributeError, TypeError):
+                autocast_ctx = torch.cuda.amp.autocast()
+
+            with autocast_ctx:
                 outputs = model(input_ids=input_ids, attention_mask=attention_mask)
                 loss = criterion(outputs.logits, labels)
             scaler.scale(loss).backward()
@@ -378,7 +386,10 @@ def main():
     scheduler = get_linear_schedule_with_warmup(
         optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps
     )
-    scaler = torch.amp.GradScaler("cuda", enabled=(device.type == "cuda"))
+    try:
+        scaler = torch.amp.GradScaler("cuda", enabled=(device.type == "cuda"))
+    except (AttributeError, TypeError):
+        scaler = torch.cuda.amp.GradScaler(enabled=(device.type == "cuda"))
 
     # Training Loop with Early Stopping monitoring Macro F1
     best_val_macro_f1 = -1.0