Update open_oasis_master/generate.py

open_oasis_master/generate.py

@@ -10,8 +10,9 @@ from utils import one_hot_actions, sigmoid_beta_schedule
 from tqdm import tqdm
 from einops import rearrange
 from torch import autocast
-assert torch.cuda.is_available()
-device = "cuda:0"
+#assert torch.cuda.is_available()
+#device = "cuda:0"
+device = "cpu"
 
 # load DiT checkpoint
 ckpt = torch.load("oasis500m.pt")