fix type checking

Signed-off-by: Chen Cui <chcui@nvidia.com>
NVIDIA · Apr 26, 2024 · ac3e158 · ac3e158
1 parent a79bad6
commit ac3e158
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nemo/lightning/pytorch/strategies.py b/nemo/lightning/pytorch/strategies.py
@@ -54,7 +54,7 @@ def __init__(
         pipeline_model_parallel_size: int = 1,
         virtual_pipeline_model_parallel_size: Optional[int] = None,
         sequence_parallel: bool = False,
-        data_sampler: Optional[DataSampler] = None,
+        data_sampler: Optional['DataSampler'] = None,
         parallel_devices: Optional[List[torch.device]] = None,
         cluster_environment=None,  # TODO: Add type-hint
         checkpoint_io=None,  # TODO: Add type-hint
@@ -72,7 +72,7 @@ def __init__(
         )
         self.no_ddp_communication_hook = no_ddp_communication_hook
         self.megatron_callbacks = CallbackConnector()
-        self.data_sampler: Optional[DataSampler] = data_sampler
+        self.data_sampler: Optional['DataSampler'] = data_sampler
         self.tensor_model_parallel_size = tensor_model_parallel_size
         self.pipeline_model_parallel_size = pipeline_model_parallel_size
         self.virtual_pipeline_model_parallel_size = virtual_pipeline_model_parallel_size