rynmurdock commited on
Commit
5ec2d9d
1 Parent(s): e155b53
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -75,8 +75,8 @@ class BottleneckT5Autoencoder:
75
  def __init__(self, model_path: str, device='cuda'):
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:
 
75
  def __init__(self, model_path: str, device='cuda'):
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).to('cuda')
79
+ self.model.eval()
80
 
81
 
82
  def embed(self, text: str) -> torch.FloatTensor: