.

app.py

@@ -75,7 +75,7 @@ class BottleneckT5Autoencoder:
     def __init__(self, model_path: str, device='cuda'):
         self.device = device
         self.tokenizer = AutoTokenizer.from_pretrained(model_path, model_max_length=512, torch_dtype=torch.bfloat16)
-        self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, device=DEVICE, low_cpu_mem_usage=True)
+        self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, low_cpu_mem_usage=True)
         self.model.eval()