minor fix to remove runtime error

jmkim0309 · jmkim0309 · commit 300acc1dc57c · 2023-01-04T19:24:23.000+09:00
diff --git a/karlo/modules/diffusion/respace.py b/karlo/modules/diffusion/respace.py
@@ -104,8 +104,9 @@ def condition_score(self, cond_fn, *args, **kwargs):
 
     def _wrap_model(self, model):
         def wrapped(x, ts, **kwargs):
+            ts_cpu = ts.detach().to("cpu")
             return model(
-                x, self.timestep_map[ts].to(device=ts.device, dtype=ts.dtype), **kwargs
+                x, self.timestep_map[ts_cpu].to(device=ts.device, dtype=ts.dtype), **kwargs
             )
 
         return wrapped

Original file line number	Diff line number	Diff line change
`@@ -104,8 +104,9 @@ def condition_score(self, cond_fn, args, *kwargs):`
`104`	`104`
`105`	`105`	`def _wrap_model(self, model):`
`106`	`106`	`def wrapped(x, ts, **kwargs):`
	`107`	`+ ts_cpu = ts.detach().to("cpu")`
`107`	`108`	`return model(`
`108`		`- x, self.timestep_map[ts].to(device=ts.device, dtype=ts.dtype), **kwargs`
	`109`	`+ x, self.timestep_map[ts_cpu].to(device=ts.device, dtype=ts.dtype), **kwargs`
`109`	`110`	`)`
`110`	`111`
`111`	`112`	`return wrapped`