BUG: Resolve multiple values for keyword argument 'pin_memory' (#296)

hi-ml-histopathology/src/histopathology/datamodules/base_module.py

@@ -213,7 +213,6 @@ class TilesDataModule(HistoDataModule[TilesDataset]):
             collate_fn=multibag_collate,
             shuffle=shuffle,
             generator=generator,
-            pin_memory=True,
             **dataloader_kwargs,
         )
 
@@ -313,7 +312,6 @@ class SlidesDataModule(HistoDataModule[SlidesDataset]):
             collate_fn=image_collate,
             shuffle=shuffle,
             generator=generator,
-            pin_memory=True,
             **dataloader_kwargs,
         )
 

hi-ml-histopathology/src/histopathology/utils/layer_utils.py

@@ -15,7 +15,7 @@ def get_imagenet_preprocessing() -> nn.Module:
 
 
 def setup_feature_extractor(pretrained_model: nn.Module,
-                            input_dim: Tuple[int, int, int]) -> Tuple[nn.Module, int]:
+                            input_dim: Tuple[int, ...]) -> Tuple[nn.Module, int]:
     try:
         # Attempt to auto-detect final classification layer:
         num_features: int = pretrained_model.fc.in_features  # type: ignore