Fix: sensor_data_to_rgbd checks if 1 env, but not if force_gpu is true

mani_skill2/envs/utils/observations/observations.py

@@ -7,6 +7,7 @@
 
 import numpy as np
 import torch
+import sapien.physx as physx
 
 from mani_skill2.sensors.base_sensor import BaseSensor, BaseSensorConfig
 from mani_skill2.sensors.camera import Camera
@@ -30,7 +31,7 @@ def sensor_data_to_rgbd(observation: Dict, sensors: Dict[str, BaseSensor]):
  depth = -ori_images[key][..., [2]] # [H, W, 1]
  # NOTE (stao): This is a bit of a hack since normally we have generic to_numpy call to convert internal torch tensors to numpy if we do not use GPU simulation
  # but torch does not have a uint16 type so we convert that here earlier
- if depth.shape[0] == 1:
+ if not physx.is_gpu_enabled and depth.shape[0] == 1:
  depth = depth.numpy().astype(np.uint16)
  new_images["depth"] = depth
  else: