rynmurdock commited on
Commit
e155b53
·
1 Parent(s): f69254d
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -76,7 +76,7 @@ class BottleneckT5Autoencoder:
76
  self.device = device
77
  self.tokenizer = AutoTokenizer.from_pretrained(model_path, model_max_length=512, torch_dtype=torch.bfloat16)
78
  self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, low_cpu_mem_usage=True)
79
- self.model.eval()
80
 
81
 
82
  def embed(self, text: str) -> torch.FloatTensor:
 
76
  self.device = device
77
  self.tokenizer = AutoTokenizer.from_pretrained(model_path, model_max_length=512, torch_dtype=torch.bfloat16)
78
  self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, low_cpu_mem_usage=True)
79
+ self.model.eval().to('cuda')
80
 
81
 
82
  def embed(self, text: str) -> torch.FloatTensor: