tree3po commited on
Commit
e1993f3
1 Parent(s): 6f1f855

Update open_oasis_master/generate.py

Browse files
Files changed (1) hide show
  1. open_oasis_master/generate.py +1 -1
open_oasis_master/generate.py CHANGED
@@ -97,7 +97,7 @@ for i in tqdm(range(n_prompt_frames, total_frames)):
97
 
98
  # get model predictions
99
  with torch.no_grad():
100
- with autocast("cuda", dtype=torch.half):
101
  v = model(x_curr, t, actions[:, start_frame : i + 1])
102
 
103
  x_start = alphas_cumprod[t].sqrt() * x_curr - (1 - alphas_cumprod[t]).sqrt() * v
 
97
 
98
  # get model predictions
99
  with torch.no_grad():
100
+ with autocast("cpu", dtype=torch.half):
101
  v = model(x_curr, t, actions[:, start_frame : i + 1])
102
 
103
  x_start = alphas_cumprod[t].sqrt() * x_curr - (1 - alphas_cumprod[t]).sqrt() * v