gelnesr commited on
Commit
bbe790e
·
verified ·
1 Parent(s): e640230

move to cpu

Browse files
Files changed (1) hide show
  1. Dyna-1/esm/models/esm3.py +1 -1
Dyna-1/esm/models/esm3.py CHANGED
@@ -235,7 +235,7 @@ class ESM3(nn.Module, ESM3InferenceClient):
235
  if not model_name:
236
  raise ValueError(f"Model name {model_name} is not a valid ESM3 model name.")
237
  if device is None:
238
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
239
  model = load_local_model(model_name, device=device)
240
  if device.type != "cpu":
241
  model = model.to(torch.bfloat16)
 
235
  if not model_name:
236
  raise ValueError(f"Model name {model_name} is not a valid ESM3 model name.")
237
  if device is None:
238
+ device = torch.device("cpu")
239
  model = load_local_model(model_name, device=device)
240
  if device.type != "cpu":
241
  model = model.to(torch.bfloat16)