code: def test_nll(self):
      device_mesh = DeviceMesh(
          self.device_type, list(range(self.world_size))
      )
      shard_spec = [Replicate()]
  
      pred_list = torch.rand(ITER_TIME, 1024, 10)
      target_list = torch.randint(
          0, 10, (ITER_TIME, 1024), dtype=torch.long
      )
  
      criterion = torch.nn.CrossEntropyLoss()
  
      for i in range(ITER_TIME):
          pred = pred_list[i].to(self.device_type).requires_grad_()
          target = target_list[i].to(self.device_type)
  
          # nll with dtensor
          pred_dtensor = distribute_tensor(
              pred, device_mesh, shard_spec
          )
          target_dtensor = distribute_tensor(
              target, device_mesh, shard_spec
          )
          loss = criterion(pred_dtensor, target_dtensor)
          loss.backward()
  
          # nll with plain tensor
          loss_gt = criterion(pred, target)
          loss_gt.backward()
  
          loss_diff_abs = loss.to_local() - loss_gt
          loss_diff_rel = loss_diff_abs / (torch.abs(loss_gt) + 1e-8)
          loss_mse_abs = torch.mean(
              loss_diff_abs * loss_diff_abs
          ).item()
          loss_mse_rel = torch.mean(
              loss_diff_rel * loss_diff_rel
          ).item()
  
          grad_diff_abs = pred_dtensor.grad.to_local() - pred.grad
          grad_diff_rel = grad_diff_abs / (torch.abs(pred.grad) + 1e-8)
          grad_mse_abs = torch.mean(
              grad_diff_abs * grad_diff_abs
          ).item()
          grad_mse_rel = torch.mean(
              grad_diff_rel * grad_diff_rel
          ).item()
  
          self.assertTrue(
              loss_mse_abs <= 1e-6,
              f"Too large absolute mse for loss, expected less equal 1e-6, got {loss_mse_abs}",
          )
          self.assertTrue(
              loss_mse_rel <= 1e-6,
              f"Too large relative mse for loss, expected less equal 1e-6, got {loss_mse_rel}",
          )
          self.assertTrue(
              grad_mse_abs <= 1e-6,
              f"Too large absolute mse for gradient, expected less equal 1e-6, got {grad_mse_abs}",
          )
          self.assertTrue(
              grad_mse_rel <= 1e-6,
              f"Too large relative mse for gradient, expected less equal 1e-6, got {grad_mse_rel}",
          )
  
file_path: pytorch/test/distributed/_tensor/test_experimental_ops.py
imports: import torch
  import torch.distributed
  from torch.distributed._tensor import DeviceMesh, distribute_tensor, Replicate
  from torch.testing._internal.common_utils import run_tests
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, with_comms
label_desc: Use image node_id 4 for calling the DistOtherOpsTest obj's underlying member method code with example usage: obj.test_nll() without return types