🐛 Fix tensor shape in reset

yhs0602 · Dec 31, 2024 · d21d41c · d21d41c
1 parent 643b2d6
commit d21d41c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/craftground/environment.py b/src/craftground/environment.py
@@ -440,6 +440,7 @@ def convert_observation(
             self.csv_logger.profile_end("convert_observation/decode_raw")
         elif self.encoding_mode == ScreenEncodingMode.ZEROCOPY:
             import torch
+
             if self.initial_env.eye_distance > 0:  # binocular vision
                 # TODO: Handle binocular vision
                 if (
@@ -469,6 +470,7 @@ def convert_observation(
                 print(rgb_array_or_tensor.device)
                 # print(image_tensor)
                 self.observation_tensors[0] = rgb_array_or_tensor
+                rgb_array_or_tensor = rgb_array_or_tensor.clone()[:, :, :3]
             else:
                 # TODO: Handle cuda case also
                 cuda_dl_tensor = mtl_tensor_from_cuda_mem_handle(