UT import fix (#3233)

Summary:
Pull Request resolved: #3233

as title

Reviewed By: SSYernar

Differential Revision: D78856958

fbshipit-source-id: 827a078f22e3bc59dde741a81498d70498843793

Author: aliafzal

torchrec/distributed/train_pipeline/tests/test_train_pipelines_utils.py

@@ -13,8 +13,9 @@
 from typing import List
 from unittest.mock import MagicMock
 
+import parameterized
+
 import torch
-from parameterized import parameterized
 
 from torchrec.distributed.embedding_types import EmbeddingComputeKernel
 from torchrec.distributed.test_utils.test_model import ModelInput, TestNegSamplingModule
@@ -260,7 +261,7 @@ def test_restore_from_snapshot(self) -> None:
         for source_model_type, recipient_model_type in variants:
             self._test_restore_from_snapshot(source_model_type, recipient_model_type)
 
-    @parameterized.expand(
+    @parameterized.parameterized.expand(
         [
             (
                 CallArgs(