Merge pull request #195 from allenai/favyen/20250826-fix-helios

Fix a couple bugs with Helios

Author: favyen2

helios.Dockerfile

@@ -6,8 +6,6 @@ RUN apt install -y libpq-dev ffmpeg libsm6 libxext6 git wget
 # Install rslearn and helios (need to be in local directory).
 COPY ./docker_build/rslearn /opt/rslearn
 COPY ./docker_build/helios /opt/helios
-COPY requirements.txt /opt/rslearn_projects/requirements.txt
-COPY ai2_requirements.txt /opt/rslearn_projects/ai2_requirements.txt
 
 # We also install terratorch so that we can use the same Docker image for TerraMind
 # experiments.
@@ -16,7 +14,10 @@ RUN pip install --no-cache-dir geobench==0.0.1
 
 RUN pip install --no-cache-dir --upgrade /opt/rslearn[extra]
 RUN pip install --no-cache-dir --upgrade /opt/helios
-RUN pip install --no-cache-dir -r /opt/rslearn_projects/requirements.txt -r /opt/rslearn_projects/ai2_requirements.txt
+
+COPY requirements-without-rslearn.txt /opt/rslearn_projects/requirements-without-rslearn.txt
+COPY requirements-extra.txt /opt/rslearn_projects/requirements-extra.txt
+RUN pip install --no-cache-dir -r /opt/rslearn_projects/requirements-without-rslearn.txt -r /opt/rslearn_projects/requirements-extra.txt
 
 # Copy rslearn_projects and install it too.
 COPY . /opt/rslearn_projects/

rslp/helios/model.py

@@ -162,7 +162,7 @@ def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
             # Currently we assume the provided model always returns a TokensAndMasks object.
             tokens_and_masks: TokensAndMasks = self.model(
                 sample, always_pass_none_mask_to_transformer=True, **self.forward_kwargs
-            )[0]
+            )["tokens_and_masks"]
 
         # Apply temporal/modality pooling so we just have one feature per patch.
         features = []