# 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: [chns, H, W] or [chns, 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
        align_corners: Defaults to ``self.align_corners``.
            See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample

    Raises:
        ValueError: When ``img`` spatially is not one of [2D, 3D].

    """
    im_shape = np.asarray(img.shape[1:])  # spatial dimensions
    input_ndim = len(im_shape)
    if input_ndim not in (2, 3):
        raise ValueError(
            f"Unsupported img dimension: {input_ndim}, available options are [2, 3]."
        )
    _angle = ensure_tuple_rep(self.angle, 1 if input_ndim == 2 else 3)
    _rad = np.deg2rad(_angle)
    transform = create_rotate(input_ndim, _rad)
    shift = create_translate(input_ndim, (im_shape - 1) / 2)
    if self.keep_size:
        output_shape = im_shape
    else:
        corners = np.asarray(
            np.meshgrid(*[(0, dim) for dim in im_shape], indexing="ij")
        ).reshape((len(im_shape), -1))
        corners = transform[:-1, :-1] @ corners
        output_shape = (corners.ptp(axis=1) + 0.5).astype(int)
    shift_1 = create_translate(input_ndim, -(output_shape - 1) / 2)
    transform = shift @ transform @ shift_1
    _dtype = img.dtype
    xform = AffineTransform(
        normalized=False,
        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,
        reverse_indexing=True,
    )
    output = xform(
        torch.from_numpy(img.astype(np.float64)).unsqueeze(0),
        torch.from_numpy(transform.astype(np.float64)),
        spatial_size=output_shape,
    )
    output = output.squeeze(0).detach().cpu().numpy().astype(_dtype)
    return output


# 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: [chns, H, W] or [chns, 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
        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``.

    Raises:
        ValueError: When ``img`` spatially is not one of [2D, 3D].

    """
    _dtype = dtype or self.dtype or img.dtype
    im_shape = np.asarray(img.shape[1:])  # spatial dimensions
    input_ndim = len(im_shape)
    if input_ndim not in (2, 3):
        raise ValueError(
            f"Unsupported img dimension: {input_ndim}, available options are [2, 3]."
        )
    _angle = ensure_tuple_rep(self.angle, 1 if input_ndim == 2 else 3)
    _rad = np.deg2rad(_angle)
    transform = create_rotate(input_ndim, _rad)
    shift = create_translate(input_ndim, (im_shape - 1) / 2)
    if self.keep_size:
        output_shape = im_shape
    else:
        corners = np.asarray(
            np.meshgrid(*[(0, dim) for dim in im_shape], indexing="ij")
        ).reshape((len(im_shape), -1))
        corners = transform[:-1, :-1] @ corners
        output_shape = (corners.ptp(axis=1) + 0.5).astype(int)
    shift_1 = create_translate(input_ndim, -(output_shape - 1) / 2)
    transform = shift @ transform @ shift_1

    xform = AffineTransform(
        normalized=False,
        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,
        reverse_indexing=True,
    )
    output = xform(
        torch.as_tensor(np.ascontiguousarray(img).astype(_dtype)).unsqueeze(0),
        torch.as_tensor(np.ascontiguousarray(transform).astype(_dtype)),
        spatial_size=output_shape,
    )
    output = output.squeeze(0).detach().cpu().numpy().astype(np.float32)
    return output