Upload main.py

main.py

@@ -38,6 +38,7 @@ if not os.environ.get("TLM_DATA_DIR"):
     os.environ["TLM_DATA_DIR"] = "/tmp/tlm_data"
 
 # Select GPU if available
+import torch
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 logger.info(f"Using device: {device}")
 
@@ -47,15 +48,6 @@ try:
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
         logger.info(f"GPU Memory: {torch.cuda.memory_allocated()/1e9:.2f}GB / {torch.cuda.get_device_properties(0).total_memory/1e9:.2f}GB")
-        
-    # Apply PyTorch device fixes
-    if hasattr(torch, 'set_default_device'):
-        torch.set_default_device("cpu")
-    _original_device = torch.device
-    def patched_device(device_str=None):
-        return _original_device("cpu") if device_str is None else _original_device("cpu")
-    torch.device = patched_device
-    logger.info("✅ Applied PyTorch device attribute fix")
 except Exception as e:
     logger.warning(f"Error with PyTorch setup: {e}")