# Buggy code
def __call__(
    self,
    img: np.ndarray,
    mode: Optional[Union[GridSampleMode, str]] = None,
    padding_mode: Optional[Union[GridSamplePadMode, str]] = None,
    align_corners: Optional[bool] = None,
) -> np.ndarray:
    """
    Args:
        img: channel first array, must have shape 2D: (nchannels, H, W), or 3D: (nchannels, H, W, D).
        mode: {``"bilinear"``, ``"nearest"``}
            Interpolation mode to calculate output values. Defaults to ``self.mode``.
            See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
        padding_mode: {``"zeros"``, ``"border"``, ``"reflection"``}
            Padding mode for outside grid values. Defaults to ``self.padding_mode``.
            See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
        align_corners: Defaults to ``self.align_corners``.
            See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
    """
    self.randomize()
    if not self._do_transform:
        return img
    rotator = Rotate(
        angle=self.x if img.ndim == 3 else (self.x, self.y, self.z),
        keep_size=self.keep_size,
        mode=mode or self.mode,
        padding_mode=padding_mode or self.padding_mode,
        align_corners=self.align_corners if align_corners is None else align_corners,
    )
    return rotator(img)


# Traceback (AssertionError)
======================================================================
FAIL: test_training (__main__.IntegrationSegmentation3D)
----------------------------------------------------------------------
Traceback (most recent call last):
File "tests/test_integration_segmentation_3d.py", line 257, in test_training
rtol=1e-3,
File "/opt/conda/lib/python3.6/site-packages/numpy/testing/_private/utils.py", line 1533, in assert_allclose
verbose=verbose, header=header, equal_nan=equal_nan)
File "/opt/conda/lib/python3.6/site-packages/numpy/testing/_private/utils.py", line 846, in assert_array_compare
raise AssertionError(msg)
AssertionError:
Not equal to tolerance rtol=0.001, atol=0

Mismatched elements: 6 / 6 (100%)
Max absolute difference: 0.01720129
Max relative difference: 0.04047175
x: array([0.543151, 0.471052, 0.453605, 0.438546, 0.437794, 0.407818])
y: array([0.544673, 0.475109, 0.444963, 0.427036, 0.433381, 0.42502 ])

----------------------------------------------------------------------
Ran 1 test in 92.377s

FAILED (failures=1)

# Fixed code
def __call__(
    self,
    img: np.ndarray,
    mode: Optional[Union[GridSampleMode, str]] = None,
    padding_mode: Optional[Union[GridSamplePadMode, str]] = None,
    align_corners: Optional[bool] = None,
    dtype: Optional[np.dtype] = None,
) -> np.ndarray:
    """
    Args:
        img: channel first array, must have shape 2D: (nchannels, H, W), or 3D: (nchannels, H, W, D).
        mode: {``"bilinear"``, ``"nearest"``}
            Interpolation mode to calculate output values. Defaults to ``self.mode``.
            See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
        padding_mode: {``"zeros"``, ``"border"``, ``"reflection"``}
            Padding mode for outside grid values. Defaults to ``self.padding_mode``.
            See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
        align_corners: Defaults to ``self.align_corners``.
            See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
        dtype: data type for resampling computation. Defaults to ``self.dtype``.
            If None, use the data type of input data. To be compatible with other modules,
            the output data type is always ``np.float32``.
    """
    self.randomize()
    if not self._do_transform:
        return img
    rotator = Rotate(
        angle=self.x if img.ndim == 3 else (self.x, self.y, self.z),
        keep_size=self.keep_size,
        mode=mode or self.mode,
        padding_mode=padding_mode or self.padding_mode,
        align_corners=self.align_corners if align_corners is None else align_corners,
        dtype=dtype or self.dtype or img.dtype,
    )
    return rotator(img)