[refactor] Address Shuhei's comments

Author: nabenabe0928

autoPyTorch/pipeline/components/training/trainer/GridCutMixTrainer.py

@@ -26,14 +26,15 @@ def data_preparation(self, X: np.ndarray, y: np.ndarray,
             np.ndarray: that processes data
             typing.Dict[str, np.ndarray]: arguments to the criterion function
         """
-        beta = 1.0
-        lam = self.random_state.beta(beta, beta)
-        batch_size, channel, W, H = X.size()
-        index = torch.randperm(batch_size).cuda() if X.is_cuda else torch.randperm(batch_size)
+        alpha, beta = 1.0, 1.0
+        lam = self.random_state.beta(alpha, beta)
+        batch_size, _, W, H = X.shape
+        device = torch.device('cuda' if X.is_cuda else 'cpu')
+        batch_indices = torch.randperm(batch_size).to(device)
 
         r = self.random_state.rand(1)
         if beta <= 0 or r > self.alpha:
-            return X, {'y_a': y, 'y_b': y[index], 'lam': 1}
+            return X, {'y_a': y, 'y_b': y[batch_indices], 'lam': 1}
 
         # Draw parameters of a random bounding box
         # Where to cut basically
@@ -47,12 +48,13 @@ def data_preparation(self, X: np.ndarray, y: np.ndarray,
         bbx2 = np.clip(cx + cut_w // 2, 0, W)
         bby2 = np.clip(cy + cut_h // 2, 0, H)
 
-        X[:, :, bbx1:bbx2, bby1:bby2] = X[index, :, bbx1:bbx2, bby1:bby2]
+        X[:, :, bbx1:bbx2, bby1:bby2] = X[batch_indices, :, bbx1:bbx2, bby1:bby2]
 
         # Adjust lam
-        lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (X.size()[-1] * X.size()[-2]))
+        pixel_size = W * H
+        lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / pixel_size)
 
-        y_a, y_b = y, y[index]
+        y_a, y_b = y, y[batch_indices]
 
         return X, {'y_a': y_a, 'y_b': y_b, 'lam': lam}