nits

Author: datumbox

torchvision/prototype/transforms/_auto_augment.py

@@ -278,7 +278,7 @@ def forward(self, *inputs: Any) -> Any:
         sample = inputs if len(inputs) > 1 else inputs[0]
 
         image = query_image(sample)
-        _, height, width = get_image_dims(image)
+        _, *image_size = get_image_dims(image)
 
         policy = self._policies[int(torch.randint(len(self._policies), ()))]
 
@@ -288,7 +288,7 @@ def forward(self, *inputs: Any) -> Any:
 
             magnitudes_fn, signed = self._AUGMENTATION_SPACE[transform_id]
 
-            magnitudes = magnitudes_fn(10, (height, width))
+            magnitudes = magnitudes_fn(10, image_size)
             if magnitudes is not None:
                 magnitude = float(magnitudes[magnitude_idx])
                 if signed and torch.rand(()) <= 0.5:
@@ -334,12 +334,12 @@ def forward(self, *inputs: Any) -> Any:
         sample = inputs if len(inputs) > 1 else inputs[0]
 
         image = query_image(sample)
-        _, height, width = get_image_dims(image)
+        _, *image_size = get_image_dims(image)
 
         for _ in range(self.num_ops):
             transform_id, (magnitudes_fn, signed) = self._get_random_item(self._AUGMENTATION_SPACE)
 
-            magnitudes = magnitudes_fn(self.num_magnitude_bins, (height, width))
+            magnitudes = magnitudes_fn(self.num_magnitude_bins, image_size)
             if magnitudes is not None:
                 magnitude = float(magnitudes[int(torch.randint(self.num_magnitude_bins, ()))])
                 if signed and torch.rand(()) <= 0.5:
@@ -383,11 +383,11 @@ def forward(self, *inputs: Any) -> Any:
         sample = inputs if len(inputs) > 1 else inputs[0]
 
         image = query_image(sample)
-        _, height, width = get_image_dims(image)
+        _, *image_size = get_image_dims(image)
 
         transform_id, (magnitudes_fn, signed) = self._get_random_item(self._AUGMENTATION_SPACE)
 
-        magnitudes = magnitudes_fn(self.num_magnitude_bins, (height, width))
+        magnitudes = magnitudes_fn(self.num_magnitude_bins, image_size)
         if magnitudes is not None:
             magnitude = float(magnitudes[int(torch.randint(self.num_magnitude_bins, ()))])
             if signed and torch.rand(()) <= 0.5: