From 2323ded0f867a4effb10d1c67c1ca3f730fc689c Mon Sep 17 00:00:00 2001 From: David Hoeller Date: Wed, 30 Oct 2024 11:22:19 +0100 Subject: [PATCH] Fixes image feature extractor setup --- .../omni.isaac.lab/omni/isaac/lab/envs/mdp/observations.py | 2 +- source/extensions/omni.isaac.lab/setup.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/observations.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/observations.py index aca0f579ce..cfaac0a653 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/observations.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/observations.py @@ -381,7 +381,7 @@ def __call__( print(f"[INFO]: Adding {model_name} to the model zoo") self.model_zoo[model_name] = self.model_zoo_cfg[model_name]["model"]() - if model_device is not None and self.model_zoo[model_name].device != model_device: + if model_device is not None: # want to offload vision model inference to another device self.model_zoo[model_name] = self.model_zoo[model_name].to(model_device) diff --git a/source/extensions/omni.isaac.lab/setup.py b/source/extensions/omni.isaac.lab/setup.py index 3b68d3c84f..01229c8c4f 100644 --- a/source/extensions/omni.isaac.lab/setup.py +++ b/source/extensions/omni.isaac.lab/setup.py @@ -31,6 +31,9 @@ # procedural-generation "trimesh", "pyglet<2", + # image processing + "transformers", + "einops", # needed for transformers, doesn't always auto-install ] PYTORCH_INDEX_URL = ["https://download.pytorch.org/whl/cu118"]