Add SingleDeviceDistributedConfig to imports (#241)

Summary: Pull Request resolved: #241

Reviewed By: wz337

Differential Revision: D81140356

Pulled By: tsunghsienlee

fbshipit-source-id: c79967015a350a7474514920cfa5e466ff3db434

Author: runame

distributed_shampoo/__init__.py

@@ -58,6 +58,7 @@
     ShampooPreconditionerConfig,
     ShampooPT2CompileConfig,
     SignDescentPreconditionerConfig,
+    SingleDeviceDistributedConfig,
     SpectralDescentPreconditionerConfig,
 )
 
@@ -73,8 +74,9 @@
     # PT2 compile.
     "ShampooPT2CompileConfig",
     # `distributed_config` options.
-    "DefaultSingleDeviceDistributedConfig",
     "DistributedConfig",  # Abstract base class.
+    "SingleDeviceDistributedConfig",
+    "DefaultSingleDeviceDistributedConfig",  # Default `SingleDeviceDistributedConfig`.
     "DDPDistributedConfig",
     "FSDPParamAssignmentStrategy",
     "FSDPDistributedConfig",