def METHOD_NAME(self):
    super().METHOD_NAME()
    # Make sure all async ops finish.
    context.async_wait()
    # TODO(hthu): Remove the reset once we fixed the CopyToMesh with
    # DefaultMesh placement issue.
    reset_dtensor()