# Buggy code
def __init__(
    self,
    keys: KeysCollection,
    range_x: Union[Tuple[float, float], float] = 0.0,
    range_y: Union[Tuple[float, float], float] = 0.0,
    range_z: Union[Tuple[float, float], float] = 0.0,
    prob: float = 0.1,
    keep_size: bool = True,
    mode: GridSampleModeSequence = GridSampleMode.BILINEAR,
    padding_mode: GridSamplePadModeSequence = GridSamplePadMode.BORDER,
    align_corners: Union[Sequence[bool], bool] = False,
) -> None:
    super().__init__(keys)
    self.range_x = ensure_tuple(range_x)
    if len(self.range_x) == 1:
        self.range_x = tuple(sorted([-self.range_x[0], self.range_x[0]]))
    self.range_y = ensure_tuple(range_y)
    if len(self.range_y) == 1:
        self.range_y = tuple(sorted([-self.range_y[0], self.range_y[0]]))
    self.range_z = ensure_tuple(range_z)
    if len(self.range_z) == 1:
        self.range_z = tuple(sorted([-self.range_z[0], self.range_z[0]]))

    self.prob = prob
    self.keep_size = keep_size
    self.mode = ensure_tuple_rep(mode, len(self.keys))
    self.padding_mode = ensure_tuple_rep(padding_mode, len(self.keys))
    self.align_corners = ensure_tuple_rep(align_corners, len(self.keys))

    self._do_transform = False
    self.x = 0.0
    self.y = 0.0
    self.z = 0.0


# 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 __init__(
    self,
    keys: KeysCollection,
    range_x: Union[Tuple[float, float], float] = 0.0,
    range_y: Union[Tuple[float, float], float] = 0.0,
    range_z: Union[Tuple[float, float], float] = 0.0,
    prob: float = 0.1,
    keep_size: bool = True,
    mode: GridSampleModeSequence = GridSampleMode.BILINEAR,
    padding_mode: GridSamplePadModeSequence = GridSamplePadMode.BORDER,
    align_corners: Union[Sequence[bool], bool] = False,
    dtype: Union[Sequence[Optional[np.dtype]], Optional[np.dtype]] = np.float64,
) -> None:
    super().__init__(keys)
    self.range_x = ensure_tuple(range_x)
    if len(self.range_x) == 1:
        self.range_x = tuple(sorted([-self.range_x[0], self.range_x[0]]))
    self.range_y = ensure_tuple(range_y)
    if len(self.range_y) == 1:
        self.range_y = tuple(sorted([-self.range_y[0], self.range_y[0]]))
    self.range_z = ensure_tuple(range_z)
    if len(self.range_z) == 1:
        self.range_z = tuple(sorted([-self.range_z[0], self.range_z[0]]))

    self.prob = prob
    self.keep_size = keep_size
    self.mode = ensure_tuple_rep(mode, len(self.keys))
    self.padding_mode = ensure_tuple_rep(padding_mode, len(self.keys))
    self.align_corners = ensure_tuple_rep(align_corners, len(self.keys))
    self.dtype = ensure_tuple_rep(dtype, len(self.keys))

    self._do_transform = False
    self.x = 0.0
    self.y = 0.0
    self.z = 0.0