Clean up tests for tracing (#3381)

Summary:
Pull Request resolved: #3381

Clean up the unused dependencies in the unit tests of tracing.

Also rename the file to `test_tracing.py`, as the upper directory already indicates this is train pipeline tests.

Reviewed By: TroyGarden

Differential Revision: D82658018

fbshipit-source-id: 3c22d6c08eef3f17c0653bf1c24669a3a85196fd

Author: Shuangping Liu

torchrec/distributed/train_pipeline/tests/test_train_pipeline_tracing.py renamed to torchrec/distributed/train_pipeline/tests/test_tracing.py

@@ -8,26 +8,19 @@
 # pyre-strict
 
 import unittest
-from typing import List, Optional
+from typing import List
 from unittest.mock import MagicMock
 
 import parameterized
 
 import torch
-from torch import nn
-
 from torchrec.distributed.train_pipeline.pipeline_context import TrainPipelineContext
-
 from torchrec.distributed.train_pipeline.tracing import (
-    _get_leaf_module_names,
     ArgInfo,
     ArgInfoStepFactory,
     CallArgs,
     NodeArgsHelper,
-    PipelinedPostproc,
-    Tracer,
 )
-from torchrec.distributed.types import NullShardedModuleContext, ShardedModule
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
 
 

torchrec/distributed/train_pipeline/tracing.py

@@ -7,7 +7,6 @@
 
 # pyre-strict
 import logging
-from dataclasses import dataclass
 from itertools import chain
 from typing import Any, Dict, List, Optional, Set, Tuple, Union