code: def test_index_mul_half(self):
      out = index_mul_2d(
          self.input1_half, self.input2_half, self.index1
      )
      energy = (out.float() ** 2).sum() / out.numel()
      force = torch.autograd.grad(
          energy,
          self.input1_half,
          grad_outputs=torch.ones_like(energy),
          create_graph=True,
      )[0]
      loss = (out.float() ** 2).sum() / out.numel() + (
          force.float() ** 2
      ).sum()
      loss.backward()
  
      out_ = self.input1_half_[self.index1] * self.input2_half_
      energy_ = (out_.float() ** 2).sum() / out.numel()
      force_ = torch.autograd.grad(
          energy_,
          self.input1_half_,
          grad_outputs=torch.ones_like(energy),
          create_graph=True,
      )[0]
      loss = (out_.float() ** 2).sum() / out_.numel() + (
          force_.float() ** 2
      ).sum()
      loss.backward()
  
      torch.testing.assert_close(
          self.input1_half,
          self.input1_half_,
          atol=1e-3,
          rtol=1e-3,
          equal_nan=True,
      )
      torch.testing.assert_close(
          self.input2_half,
          self.input2_half_,
          atol=1e-3,
          rtol=1e-3,
          equal_nan=True,
      )
      torch.testing.assert_close(
          self.input1_half.grad,
          self.input1_half_.grad,
          atol=2e-3,
          rtol=5e-2,
          equal_nan=True,
      )
      torch.testing.assert_close(
          self.input2_half.grad,
          self.input2_half_.grad,
          atol=1e-3,
          rtol=1e-3,
          equal_nan=True,
      )
  
file_path: apex/apex/contrib/test/index_mul_2d/test_index_mul_2d.py
imports: import random
  import unittest
  import torch
label_desc: Use image node_id 3 for calling the IndexMul2dTest obj's underlying member method code with example usage: obj.test_index_mul_half() without return types