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: