fixing test setups

Author: azrael417

docker/Dockerfile

@@ -75,8 +75,8 @@ RUN pip install --ignore-installed "git+https://github.com/NVIDIA/mlperf-common.
 # torch-harmonics
 ENV FORCE_CUDA_EXTENSION 1
 ENV TORCH_CUDA_ARCH_LIST "8.0 8.6 9.0 10.0 12.0+PTX"
-ENV HARMONICS_VERSION 0.8.0
-RUN cd /opt && git clone -b v0.8.0 https://github.com/NVIDIA/torch-harmonics.git && \
+ENV HARMONICS_VERSION 0.8.3
+RUN cd /opt && git clone -b tkurth/distributed-attention https://github.com/NVIDIA/torch-harmonics.git && \
     cd torch-harmonics && \
     pip install --no-build-isolation -e .
 

tests/test_layers.py

@@ -35,6 +35,8 @@ def setUp(self):
         self.params.img_shape_y = 72
         self.params.img_local_shape_x = self.params.img_crop_shape_x = self.params.img_shape_x
         self.params.img_local_shape_y = self.params.img_crop_shape_y = self.params.img_shape_y
+        self.params.img_shape_x_resampled = self.params.img_shape_x
+        self.params.img_shape_y_resampled = self.params.img_shape_y
         self.params.img_local_offset_x = 0
         self.params.img_local_offset_y = 0
         self.params.img_crop_offset_x = 0

tests/test_losses.py

@@ -89,6 +89,8 @@ def setUp(self):
         self.params.img_local_shape_y = self.params.img_crop_shape_y = self.params.img_shape_y
         self.params.img_local_offset_x = self.params.img_crop_offset_x = 0
         self.params.img_local_offset_y = self.params.img_crop_offset_y = 0
+        self.params.img_shape_x_resampled = self.params.img_shape_x
+        self.params.img_shape_y_resampled = self.params.img_shape_y
 
         # also set the batch size for testing
         self.params.batch_size = 4

tests/test_models.py

@@ -37,6 +37,8 @@ def setUp(self):
         self.params.img_shape_y = 72
         self.params.img_local_shape_x = self.params.img_crop_shape_x = self.params.img_shape_x
         self.params.img_local_shape_y = self.params.img_crop_shape_y = self.params.img_shape_y
+        self.params.img_shape_x_resampled = self.params.img_shape_x
+        self.params.img_shape_y_resampled = self.params.img_shape_y
         self.params.img_local_offset_x = 0
         self.params.img_local_offset_y = 0
         self.params.img_crop_offset_x = 0