code: def test_cublas_baddbmm_large_input(
      self, device, batch_size, N, M, P, dtype
  ):
      cpu_dtype = dtype
      if dtype == torch.float16 or dtype == torch.bfloat16:
          cpu_dtype = torch.float32
  
      M1 = torch.rand((N, M), device=device, dtype=dtype)
      M2 = torch.rand((M, P), device=device, dtype=dtype)
      A = torch.rand((N, P), device=device, dtype=dtype)
  
      def _convert_to_cpu(t):
          return t.to(device="cpu", dtype=cpu_dtype)
  
      M1_cpu, M2_cpu, A_cpu = map(_convert_to_cpu, [M1, M2, A])
  
      # linear
      out1_cpu = torch.nn.functional.linear(
          M1_cpu, M2_cpu.t(), A_cpu
      ).to(dtype=dtype)
      out1_gpu = torch.nn.functional.linear(M1, M2.t(), A).cpu()
      self.assertEqual(out1_cpu, out1_gpu)
      # test multiply the identity matrix
      if N == M and M == P:
          M2_eye = torch.eye(N, device=device, dtype=dtype)
          out1_eye_gpu = torch.nn.functional.linear(
              M1, M2_eye.t(), torch.zeros_like(A)
          )
          self.assertEqual(M1_cpu.to(dtype=dtype), out1_eye_gpu.cpu())
  
      # baddbmm
      def _expand_to_batch(t: torch.Tensor):
          return t.expand((batch_size,) + t.size())
  
      alpha, beta = 1.0, 1.0
      M1, M2, A, M1_cpu, M2_cpu, A_cpu = map(
          _expand_to_batch, [M1, M2, A, M1_cpu, M2_cpu, A_cpu]
      )
  
      out2_cpu = torch.baddbmm(
          A_cpu, M1_cpu, M2_cpu, beta=beta, alpha=alpha
      ).to(dtype=dtype)
      out2_gpu = torch.baddbmm(A, M1, M2, beta=beta, alpha=alpha).cpu()
      self.assertEqual(out2_cpu, out2_gpu)
      # test multiply the identity matrix
      if N == M and M == P:
          M2_eye = torch.eye(N, device=device, dtype=dtype).expand(
              batch_size, N, N
          )
          out2_eye_gpu = torch.baddbmm(
              torch.zeros_like(A), M1, M2_eye, beta=beta, alpha=alpha
          )
          self.assertEqual(M1_cpu.to(dtype=dtype), out2_eye_gpu.cpu())
  
      # cross comparison
      self.assertEqual(out1_gpu, out2_gpu[0])
  
imports: import unittest
  from itertools import product
  from functools import partial
  from typing import Optional
  import torch
  from torch.quantization._quantized_conversions import pack_int4_to_int8, quantized_weight_reorder_for_mixed_dtypes_linear_cutlass
  from torch.testing import make_tensor
  from torch.testing._internal.common_cuda import SM53OrLater, _get_torch_cuda_version
  from torch.testing._internal.common_device_type import dtypes, instantiate_device_type_tests, onlyCUDA, tol, toleranceOverride
  from torch.testing._internal.common_utils import IS_ARM64, IS_JETSON, IS_WINDOWS, parametrize, run_tests, skipIfRocmVersionLessThan, TEST_WITH_ROCM, skipIfRocm, TestCase
label_desc: Use image node_id 7 for calling the TestMatmulCuda obj's underlying member method code with example usage: obj.test_cublas_baddbmm_large_input(device, batch_size, N, M, P, dtype) and returns: t, t