Update pytorch_lightning/plugins/__init__.py

carmocca · ananthsub · web-flow · commit c5174ad38df5 · 2021-08-25T19:28:02.000+02:00
Co-authored-by: ananthsub &lt;ananth.subramaniam@gmail.com&gt;
diff --git a/pytorch_lightning/plugins/__init__.py b/pytorch_lightning/plugins/__init__.py
@@ -33,7 +33,7 @@
 from pytorch_lightning.plugins.training_type.tpu_spawn import TPUSpawnPlugin
 from pytorch_lightning.plugins.training_type.training_type_plugin import TrainingTypePlugin
 
-PLUGIN = Union[TrainingTypePlugin, PrecisionPlugin, ClusterEnvironment]
+PLUGIN = Union[TrainingTypePlugin, PrecisionPlugin, ClusterEnvironment, CheckpointIO]
 PLUGIN_INPUT = Union[PLUGIN, str]
 
 __all__ = [
