# Buggy code
def random_affine_generator(
    batch_size: int,
    height: int,
    width: int,
    degrees: torch.Tensor,
    translate: Optional[torch.Tensor] = None,
    scale: Optional[torch.Tensor] = None,
    shear: Optional[torch.Tensor] = None,
    same_on_batch: bool = False,
) -> Dict[str, torch.Tensor]:
    r"""Get parameters for ``affine`` for a random affine transform.

    Args:
        batch_size (int): the tensor batch size.
        height (int) : height of the image.
        width (int): width of the image.
        degrees (tensor): Range of degrees to select from like (min, max).
        translate (tensor, optional): tuple of maximum absolute fraction for horizontal
            and vertical translations. For example translate=(a, b), then horizontal shift
            is randomly sampled in the range -img_width * a < dx < img_width * a and vertical shift is
            randomly sampled in the range -img_height * b < dy < img_height * b. Will not translate by default.
        scale (tensor, optional): scaling factor interval, e.g (a, b), then scale is
            randomly sampled from the range a <= scale <= b. Will keep original scale by default.
        shear (tensor, optional): Range of degrees to select from.
            Shear is a 2x2 tensor, a x-axis shear in (shear[0][0], shear[0][1]) and y-axis shear in
            (shear[1][0], shear[1][1]) will be applied. Will not apply shear by default.
        same_on_batch (bool): apply the same transformation across the batch. Default: False.

    Returns:
        params Dict[str, torch.Tensor]: parameters to be passed for transformation.
    """
    _common_param_check(batch_size, same_on_batch)
    _joint_range_check(degrees, "degrees")
    assert (
        isinstance(width, (int,))
        and isinstance(height, (int,))
        and width > 0
        and height > 0
    ), f"`width` and `height` must be positive integers. Got {width}, {height}."

    device, dtype = _extract_device_dtype([degrees, translate, scale, shear])
    angle = _adapted_uniform((batch_size,), degrees[0], degrees[1], same_on_batch)

    # compute tensor ranges
    if scale is not None:
        _joint_range_check(cast(torch.Tensor, scale[:2]), "scale")
        _scale = (
            _adapted_uniform((batch_size,), scale[0], scale[1], same_on_batch)
            .unsqueeze(1)
            .repeat(1, 2)
        )
        if len(_scale) == 4:
            _joint_range_check(cast(torch.Tensor, scale[2:]), "scale_y")
            _scale[:, 1] = _adapted_uniform(
                (batch_size,), scale[2], scale[3], same_on_batch
            )
    else:
        _scale = torch.ones((batch_size, 2), device=device, dtype=dtype)

    if translate is not None:
        _joint_range_check(cast(torch.Tensor, translate), "translate")
        max_dx: torch.Tensor = translate[0] * width
        max_dy: torch.Tensor = translate[1] * height
        translations = torch.stack(
            [
                _adapted_uniform((batch_size,), -max_dx, max_dx, same_on_batch),
                _adapted_uniform((batch_size,), -max_dy, max_dy, same_on_batch),
            ],
            dim=-1,
        )
    else:
        translations = torch.zeros((batch_size, 2), device=device, dtype=dtype)

    center: torch.Tensor = (
        torch.tensor([width, height], device=device, dtype=dtype).view(1, 2) / 2.0 - 0.5
    )
    center = center.expand(batch_size, -1)

    if shear is not None:
        _joint_range_check(cast(torch.Tensor, shear)[0], "shear")
        _joint_range_check(cast(torch.Tensor, shear)[1], "shear")
        sx = _adapted_uniform((batch_size,), shear[0][0], shear[0][1], same_on_batch)
        sy = _adapted_uniform((batch_size,), shear[1][0], shear[1][1], same_on_batch)
    else:
        sx = sy = torch.tensor([0] * batch_size)

    return dict(
        translations=translations,
        center=center,
        scale=_scale,
        angle=angle,
        sx=sx,
        sy=sy,
    )


# Traceback (TypeError)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-30-cf0a915755df> in <module>
----> 1 out = transform_fcn(x[0])

~/default-env/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
725             result = self._slow_forward(*input, **kwargs)
726         else:
--> 727             result = self.forward(*input, **kwargs)
728         for hook in itertools.chain(
729                 _global_forward_hooks.values(),

~/default-env/lib/python3.8/site-packages/torch/nn/modules/container.py in forward(self, input)
115     def forward(self, input):
116         for module in self:
--> 117             input = module(input)
118         return input
119

~/default-env/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
725             result = self._slow_forward(*input, **kwargs)
726         else:
--> 727             result = self.forward(*input, **kwargs)
728         for hook in itertools.chain(
729                 _global_forward_hooks.values(),

~/default-env/lib/python3.8/site-packages/kornia/augmentation/base.py in forward(self, input, params, return_transform)
196             return_transform = self.return_transform
197         if params is None:
--> 198             params = self.__forward_parameters__(batch_shape, self.p, self.p_batch, self.same_on_batch)
199         if 'batch_prob' not in params:
200             params['batch_prob'] = torch.tensor([True] * batch_shape[0])

~/default-env/lib/python3.8/site-packages/kornia/augmentation/base.py in __forward_parameters__(self, batch_shape, p, p_batch, same_on_batch)
92             batch_prob = batch_prob.repeat(batch_shape[0])
93         # selectively param gen
---> 94         return self.__selective_param_gen__(batch_shape, batch_prob)
95
96     def apply_func(self, input: torch.Tensor, params: Dict[str, torch.Tensor],

~/default-env/lib/python3.8/site-packages/kornia/augmentation/base.py in __selective_param_gen__(self, batch_shape, to_apply)
63     def __selective_param_gen__(
64             self, batch_shape: torch.Size, to_apply: torch.Tensor) -> Dict[str, torch.Tensor]:
---> 65         _params = self.generate_parameters(
66             torch.Size((int(to_apply.sum().item()), *batch_shape[1:])))
67         if _params is None:

~/default-env/lib/python3.8/site-packages/kornia/augmentation/augmentation.py in generate_parameters(self, batch_shape)
483
484     def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
--> 485         return rg.random_affine_generator(
486             batch_shape[0], batch_shape[-2], batch_shape[-1], self.degrees, self.translate, self.scale, self.shear,
487             self.same_on_batch)

~/default-env/lib/python3.8/site-packages/kornia/augmentation/random_generator/random_generator.py in random_affine_generator(batch_size, height, width, degrees, translate, scale, shear, same_on_batch)
173         _scale = _adapted_uniform((batch_size,), scale[0], scale[1], same_on_batch).unsqueeze(1).repeat(1, 2)
174         if len(_scale) == 4:
--> 175             _joint_range_check(cast(torch.Tensor, scale[2:]), "scale_y")
176             _scale[:, 1] = _adapted_uniform((batch_size,), scale[2], scale[3], same_on_batch)
177     else:

~/default-env/lib/python3.8/site-packages/kornia/augmentation/utils/param_validation.py in _joint_range_check(ranged_factor, name, bounds)
45             raise ValueError(f"{name}[0] should be smaller than {name}[1] got {ranged_factor}")
46     else:
---> 47         raise TypeError(
48             f"{name} should be a float number or a tuple with length 2 whose values between {bounds}."
49             f"Got {ranged_factor}.")

TypeError: scale_y should be a float number or a tuple with length 2 whose values between (-inf, inf).Got tensor([]).

# Fixed code
def random_affine_generator(
    batch_size: int,
    height: int,
    width: int,
    degrees: torch.Tensor,
    translate: Optional[torch.Tensor] = None,
    scale: Optional[torch.Tensor] = None,
    shear: Optional[torch.Tensor] = None,
    same_on_batch: bool = False,
) -> Dict[str, torch.Tensor]:
    r"""Get parameters for ``affine`` for a random affine transform.

    Args:
        batch_size (int): the tensor batch size.
        height (int) : height of the image.
        width (int): width of the image.
        degrees (tensor): Range of degrees to select from like (min, max).
        translate (tensor, optional): tuple of maximum absolute fraction for horizontal
            and vertical translations. For example translate=(a, b), then horizontal shift
            is randomly sampled in the range -img_width * a < dx < img_width * a and vertical shift is
            randomly sampled in the range -img_height * b < dy < img_height * b. Will not translate by default.
        scale (tensor, optional): scaling factor interval, e.g (a, b), then scale is
            randomly sampled from the range a <= scale <= b. Will keep original scale by default.
        shear (tensor, optional): Range of degrees to select from.
            Shear is a 2x2 tensor, a x-axis shear in (shear[0][0], shear[0][1]) and y-axis shear in
            (shear[1][0], shear[1][1]) will be applied. Will not apply shear by default.
        same_on_batch (bool): apply the same transformation across the batch. Default: False.

    Returns:
        params Dict[str, torch.Tensor]: parameters to be passed for transformation.
    """
    _common_param_check(batch_size, same_on_batch)
    _joint_range_check(degrees, "degrees")
    assert (
        isinstance(width, (int,))
        and isinstance(height, (int,))
        and width > 0
        and height > 0
    ), f"`width` and `height` must be positive integers. Got {width}, {height}."

    device, dtype = _extract_device_dtype([degrees, translate, scale, shear])
    angle = _adapted_uniform((batch_size,), degrees[0], degrees[1], same_on_batch)

    # compute tensor ranges
    if scale is not None:
        _joint_range_check(cast(torch.Tensor, scale[:2]), "scale")
        _scale = (
            _adapted_uniform((batch_size,), scale[0], scale[1], same_on_batch)
            .unsqueeze(1)
            .repeat(1, 2)
        )
        if len(scale) == 4:
            _joint_range_check(cast(torch.Tensor, scale[2:]), "scale_y")
            _scale[:, 1] = _adapted_uniform(
                (batch_size,), scale[2], scale[3], same_on_batch
            )
    else:
        _scale = torch.ones((batch_size, 2), device=device, dtype=dtype)

    if translate is not None:
        _joint_range_check(cast(torch.Tensor, translate), "translate")
        max_dx: torch.Tensor = translate[0] * width
        max_dy: torch.Tensor = translate[1] * height
        translations = torch.stack(
            [
                _adapted_uniform((batch_size,), -max_dx, max_dx, same_on_batch),
                _adapted_uniform((batch_size,), -max_dy, max_dy, same_on_batch),
            ],
            dim=-1,
        )
    else:
        translations = torch.zeros((batch_size, 2), device=device, dtype=dtype)

    center: torch.Tensor = (
        torch.tensor([width, height], device=device, dtype=dtype).view(1, 2) / 2.0 - 0.5
    )
    center = center.expand(batch_size, -1)

    if shear is not None:
        _joint_range_check(cast(torch.Tensor, shear)[0], "shear")
        _joint_range_check(cast(torch.Tensor, shear)[1], "shear")
        sx = _adapted_uniform((batch_size,), shear[0][0], shear[0][1], same_on_batch)
        sy = _adapted_uniform((batch_size,), shear[1][0], shear[1][1], same_on_batch)
    else:
        sx = sy = torch.tensor([0] * batch_size)

    return dict(
        translations=translations,
        center=center,
        scale=_scale,
        angle=angle,
        sx=sx,
        sy=sy,
    )