test

app.py

@@ -965,6 +965,21 @@ class DynamicsVisualizer:
             root / 'log/gs/temp/gs_pred.splat',
             rot_rev=True,
         )
+        
+        for k, v in self.preprocess_metadata.items():
+            self.preprocess_metadata[k] = v.detach().cpu() if isinstance(v, torch.Tensor) else v
+        for k, v in self.state.items():
+            self.state[k] = v.detach().cpu() if isinstance(v, torch.Tensor) else v
+        for k, v in self.params.items():
+            if isinstance(v, dict):
+                for k2, v2 in v.items():
+                    self.params[k][k2] = v2.detach().cpu() if isinstance(v2, torch.Tensor) else v2
+            else:
+                self.params[k] = v.detach().cpu() if isinstance(v, torch.Tensor) else v
+        for k, v in self.table_params.items():
+            self.table_params[k] = v.detach().cpu() if isinstance(v, torch.Tensor) else v
+        for k, v in self.gripper_params.items():
+            self.gripper_params[k] = v.detach().cpu() if isinstance(v, torch.Tensor) else v
 
         form_video = gr.Video(
             label='Predicted video',
@@ -993,6 +1008,21 @@ class DynamicsVisualizer:
         self.wp_device = wp_devices[0]
         self.torch_device = torch_devices[0]
         os.system('rm -rf ' + str(root / 'log/temp/*'))
+        
+        for k, v in self.preprocess_metadata.items():
+            self.preprocess_metadata[k] = v.to(self.torch_device) if isinstance(v, torch.Tensor) else v
+        for k, v in self.state.items():
+            self.state[k] = v.to(self.torch_device) if isinstance(v, torch.Tensor) else v
+        for k, v in self.params.items():
+            if isinstance(v, dict):
+                for k2, v2 in v.items():
+                    self.params[k][k2] = v2.to(self.torch_device) if isinstance(v2, torch.Tensor) else v2
+            else:
+                self.params[k] = v.to(self.torch_device) if isinstance(v, torch.Tensor) else v
+        for k, v in self.table_params.items():
+            self.table_params[k] = v.to(self.torch_device) if isinstance(v, torch.Tensor) else v
+        for k, v in self.gripper_params.items():
+            self.gripper_params[k] = v.to(self.torch_device) if isinstance(v, torch.Tensor) else v
 
         # im_list = []
         for i in range(15):
@@ -1045,6 +1075,21 @@ class DynamicsVisualizer:
         self.state['x_his'] = self.state['x'][None].repeat(self.cfg.sim.n_history, 1, 1)
         self.state['v_his'] *= 0.0
         self.state['v_pred'] *= 0.0
+        
+        for k, v in self.preprocess_metadata.items():
+            self.preprocess_metadata[k] = v.detach().cpu() if isinstance(v, torch.Tensor) else v
+        for k, v in self.state.items():
+            self.state[k] = v.detach().cpu() if isinstance(v, torch.Tensor) else v
+        for k, v in self.params.items():
+            if isinstance(v, dict):
+                for k2, v2 in v.items():
+                    self.params[k][k2] = v2.detach().cpu() if isinstance(v2, torch.Tensor) else v2
+            else:
+                self.params[k] = v.detach().cpu() if isinstance(v, torch.Tensor) else v
+        for k, v in self.table_params.items():
+            self.table_params[k] = v.detach().cpu() if isinstance(v, torch.Tensor) else v
+        for k, v in self.gripper_params.items():
+            self.gripper_params[k] = v.detach().cpu() if isinstance(v, torch.Tensor) else v
 
         make_video(root / 'log/temp', root / f'log/gs/temp/form_video.mp4', '%04d.png', 5)