code: def test_avoidoom():
      tensor = torch.from_numpy(np.random.random((20, 20)))
      if torch.cuda.is_available():
          tensor = tensor.cuda()
          # get default result
          default_result = torch.mm(tensor, tensor.transpose(1, 0))
  
          # when not occurred OOM error
          AvoidCudaOOM = AvoidOOM()
          result = AvoidCudaOOM.retry_if_cuda_oom(torch.mm)(
              tensor, tensor.transpose(1, 0)
          )
          assert (
              default_result.device == result.device
              and default_result.dtype == result.dtype
              and torch.equal(default_result, result)
          )
  
          # calculate with fp16 and convert back to source type
          AvoidCudaOOM = AvoidOOM(test=True)
          result = AvoidCudaOOM.retry_if_cuda_oom(torch.mm)(
              tensor, tensor.transpose(1, 0)
          )
          assert (
              default_result.device == result.device
              and default_result.dtype == result.dtype
              and torch.allclose(default_result, result, 1e-3)
          )
  
          # calculate on cpu and convert back to source device
          AvoidCudaOOM = AvoidOOM(test=True)
          result = AvoidCudaOOM.retry_if_cuda_oom(torch.mm)(
              tensor, tensor.transpose(1, 0)
          )
          assert (
              result.dtype == default_result.dtype
              and result.device == default_result.device
              and torch.allclose(default_result, result)
          )
  
          # do not calculate on cpu and the outputs will be same as input
          AvoidCudaOOM = AvoidOOM(test=True, to_cpu=False)
          result = AvoidCudaOOM.retry_if_cuda_oom(torch.mm)(
              tensor, tensor.transpose(1, 0)
          )
          assert (
              result.dtype == default_result.dtype
              and result.device == default_result.device
          )
  
      else:
          default_result = torch.mm(tensor, tensor.transpose(1, 0))
          AvoidCudaOOM = AvoidOOM()
          result = AvoidCudaOOM.retry_if_cuda_oom(torch.mm)(
              tensor, tensor.transpose(1, 0)
          )
          assert (
              default_result.device == result.device
              and default_result.dtype == result.dtype
              and torch.equal(default_result, result)
          )
  
imports: import numpy
  import pytest
  import torch
  from mmdet.utils import AvoidOOM
  from mmdet.utils.memory import cast_tensor_type
label_desc: Use image node_id 1 for calling a global function with example usage: test_avoidoom() without return types