Fix mask type (convert to bool) and handle CPU fallback

app.py

@@ -94,7 +94,7 @@ def train_model(batch_size, learning_rate, num_epochs):
     model = model.to(device)
     
     optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01)
-    scaler = torch.amp.GradScaler(enabled=True)
+    scaler = torch.amp.GradScaler(enabled=torch.cuda.is_available())
     
     scheduler = CurriculumScheduler()
     logs = []
@@ -116,9 +116,12 @@ def train_model(batch_size, learning_rate, num_epochs):
         train_loss = 0
         for batch in train_loader:
             batch = {k: v.to(device) for k, v in batch.items()}
+            mask = batch.get("mask")
+            if mask is not None:
+                mask = mask.bool()
             
-            with torch.amp.autocast(device_type='cuda', dtype=torch.float16):
-                outputs = model(batch["features"], mask=batch.get("mask"), active_components=stage_config.get("components"))
+            with torch.amp.autocast(device_type='cuda', dtype=torch.float16, enabled=torch.cuda.is_available()):
+                outputs = model(batch["features"], mask=mask, active_components=stage_config.get("components"))
                 loss, _ = compute_total_loss(outputs, {"labels": batch["labels"]}, stage_config["losses"], global_step)
             
             scaler.scale(loss).backward()
@@ -140,7 +143,10 @@ def train_model(batch_size, learning_rate, num_epochs):
         with torch.no_grad():
             for batch in val_loader:
                 batch = {k: v.to(device) for k, v in batch.items()}
-                outputs = model(batch["features"], mask=batch.get("mask"), active_components=stage_config.get("components"))
+                mask = batch.get("mask")
+                if mask is not None:
+                    mask = mask.bool()
+                outputs = model(batch["features"], mask=mask, active_components=stage_config.get("components"))
                 loss, _ = compute_total_loss(outputs, {"labels": batch["labels"]}, stage_config["losses"])
                 val_loss += loss.item()
                 if "predicted_class" in outputs: