yangliz5 commited on
Commit
64e8c41
1 Parent(s): be3fb72

fix: Update accelerator logic in predict function

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -63,7 +63,7 @@ def predict(
63
  dataloader = DataLoader(tokenized_dataset, batch_size=batch_size, num_workers=num_workers, persistent_workers=True)
64
  model = deepchopper.DeepChopper.from_pretrained("yangliz5/deepchopper")
65
 
66
- accelerator = "cpu" if torch.cuda.is_available() else "gpu"
67
  trainer = lightning.pytorch.trainer.Trainer(
68
  accelerator=accelerator,
69
  deterministic=False,
 
63
  dataloader = DataLoader(tokenized_dataset, batch_size=batch_size, num_workers=num_workers, persistent_workers=True)
64
  model = deepchopper.DeepChopper.from_pretrained("yangliz5/deepchopper")
65
 
66
+ accelerator = "gpu" if torch.cuda.is_available() else "cpu"
67
  trainer = lightning.pytorch.trainer.Trainer(
68
  accelerator=accelerator,
69
  deterministic=False,