diff --git a/torchtitan_neuron/__init__.py b/torchtitan_neuron/__init__.py index 3246efdf889..c9a9f5d7a1a 100644 --- a/torchtitan_neuron/__init__.py +++ b/torchtitan_neuron/__init__.py @@ -9,3 +9,7 @@ core torchtitan. Neuron recipes select them by converting the stock model config (for example ``to_nki_moe_config`` for each model's NKI MoE). """ + +# Import torch before spmd_types: importing spmd_types first makes the torch_neuronx +# autoload re-import the partially initialized spmd_types (circular import). +import torch # noqa: F401, E402