code: def test_autocast(self, batch_prob, box, device, dtype):
      if not hasattr(torch, "autocast"):
          pytest.skip("PyTorch version without autocast support")
  
      def mock_forward_parameters_sequential(
          batch_shape, cls, batch_prob
      ):
          named_modules = cls.get_forward_sequence()
          params = []
          for name, module in named_modules:
              if isinstance(
                  module,
                  (
                      K.base._AugmentationBase,
                      K.MixAugmentationBaseV2,
                      K.ImageSequential,
                  ),
              ):
                  with patch.object(
                      module,
                      "__batch_prob_generator__",
                      return_value=batch_prob,
                  ):
                      mod_param = module.forward_parameters(batch_shape)
  
                  param = ParamItem(name, mod_param)
              else:
                  param = ParamItem(name, None)
              batch_shape = K.container.image._get_new_batch_shape(
                  param, batch_shape
              )
              params.append(param)
          return params
  
      tfs = (
          K.RandomAffine(0.5, (0.1, 0.5), (0.5, 1.5), 1.2, p=1.0),
          K.RandomGaussianBlur((3, 3), (0.1, 3), p=1),
      )
      data_keys = ["input", "mask", box, "keypoints"]
      aug = K.AugmentationSequential(
          *tfs, data_keys=data_keys, random_apply=True
      )
      bs = len(batch_prob)
      imgs = torch.rand(bs, 3, 7, 4, dtype=dtype, device=device)
      if box == "bbox":
          bb = torch.tensor(
              [[[1.0, 1.0], [2.0, 1.0], [2.0, 2.0], [1.0, 2.0]]],
              dtype=dtype,
              device=device,
          ).expand(bs, 1, -1, -1)
      else:
          bb = torch.rand(bs, 1, 4, dtype=dtype, device=device)
  
      msk = torch.zeros_like(imgs)
      msk[..., 3:, 2] = 1.0
      points = torch.rand(bs, 1, 2, dtype=dtype, device=device)
  
      to_apply = torch.tensor(batch_prob, device=device)
  
      fwd_params = partial(
          mock_forward_parameters_sequential,
          cls=aug,
          batch_prob=to_apply,
      )
      with patch.object(aug, "forward_parameters", fwd_params):
          params = aug.forward_parameters(imgs.shape)
  
      with torch.autocast(device.type):
          outputs = aug(imgs, msk, bb, points, params=params)
  
      assert (
          outputs[0].dtype == dtype
      ), "Output image dtype should match the input dtype"
      assert (
          outputs[1].dtype == dtype
      ), "Output mask dtype should match the input dtype"
      assert (
          outputs[2].dtype == dtype
      ), "Output box dtype should match the input dtype"
      assert (
          outputs[3].dtype == dtype
      ), "Output keypoints dtype should match the input dtype"
  
imports: from functools import partial
  from unittest.mock import patch
  import pytest
  import torch
  import kornia
  import kornia.augmentation
  from kornia.augmentation.container.base import ParamItem
  from kornia.constants import BorderType
  from kornia.geometry.bbox import bbox_to_mask
  from kornia.testing import assert_close
label_desc: Use image node_id 19 for calling the TestAugmentationSequential obj's underlying member method code with example usage: obj.test_autocast(batch_prob, box, device, dtype) and returns: params