Update modeling_xvla.py

modeling_xvla.py

@@ -196,15 +196,14 @@ class XVLA(PreTrainedModel):
         enc = self.forward_vlm(input_ids, image_input, image_mask)
 
         B = input_ids.shape[0]
-        device = input_ids.device
         D = self.action_space.dim_action
 
-        x1 = torch.randn(B, self.num_actions, D, device=device)
+        x1 = torch.randn(B, self.num_actions, D, device=proprio.device, dtype=proprio.dtype)
         action = torch.zeros_like(x1)
 
         steps = max(1, int(steps))
         for i in range(steps, 0, -1):
-            t = torch.full((B,), i / steps, device=device)
+            t = torch.full((B,), i / steps, device=proprio.device, dtype=proprio.dtype)
             x_t = x1 * t.view(-1, 1, 1) + action * (1 - t).view(-1, 1, 1)
             proprio_m, x_t_m = self.action_space.preprocess(proprio, x_t)
             action = self.transformer(
@@ -277,7 +276,7 @@ class XVLA(PreTrainedModel):
 
                 # Inference
                 steps = int(payload.get("steps", 10))
-                action = self.generate_actions(**inputs, steps=steps).squeeze(0).cpu().numpy()
+                action = self.generate_actions(**inputs, steps=steps).squeeze(0).float().cpu().numpy()
                 return JSONResponse({"action": action.tolist()})
 
             except Exception: