tests/test_dtensor.py ImplicitRegistrationTest.test Failing

[sanketpurandare]

======================================================================
ERROR: test_implicit_registration (__main__.ImplicitRegistrationTest.test_implicit_registration)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/data/users/sanketpurandare/pytorch/torch/testing/_internal/common_distributed.py", line 770, in wrapper
    self._join_processes(fn)
  File "/data/users/sanketpurandare/pytorch/torch/testing/_internal/common_distributed.py", line 1038, in _join_processes
    self._check_return_codes(fn, elapsed_time)
  File "/data/users/sanketpurandare/pytorch/torch/testing/_internal/common_distributed.py", line 1078, in _check_return_codes
    raise RuntimeError(error)
RuntimeError: Process 2 exited with error code 10 and exception:
Traceback (most recent call last):
  File "/data/users/sanketpurandare/pytorch/torch/testing/_internal/common_distributed.py", line 924, in run_test
    getattr(self, test_name)()
  File "/data/users/sanketpurandare/pytorch/torch/testing/_internal/common_distributed.py", line 772, in wrapper
    fn()
  File "/data/users/sanketpurandare/pytorch/torch/testing/_internal/common_utils.py", line 3334, in wrapper
    method(*args, **kwargs)
  File "/data/users/sanketpurandare/pytorch/torch/testing/_internal/distributed/_tensor/common_dtensor.py", line 529, in wrapper
    raise e
  File "/data/users/sanketpurandare/pytorch/torch/testing/_internal/distributed/_tensor/common_dtensor.py", line 526, in wrapper
    func(self, *args, **kwargs)  # type: ignore[misc]
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/sanketpurandare/autoparallel/tests/test_dtensor.py", line 374, in test_implicit_registration
    with self.assertRaisesRegex(
         ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sanketpurandare/.conda/envs/pytorch-3.12/lib/python3.12/unittest/case.py", line 263, in __exit__
    self._raiseFailure("{} not raised".format(exc_name))
  File "/home/sanketpurandare/.conda/envs/pytorch-3.12/lib/python3.12/unittest/case.py", line 200, in _raiseFailure
    raise self.test_case.failureException(msg)
AssertionError: NotImplementedError not raised

To execute this test, run the following from the base repo dir:
    python test_dtensor.py ImplicitRegistrationTest.test_implicit_registration