Apply pre-commit formatting (isort) [skip-build]
Browse filesCo-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- test/test_muon_moe.py +2 -1
test/test_muon_moe.py
CHANGED
|
@@ -7,7 +7,8 @@ import pytest
|
|
| 7 |
import torch
|
| 8 |
import torch.distributed as dist
|
| 9 |
from optimizer.muon import Muon, get_default_muon_param_groups
|
| 10 |
-
from torch.distributed.tensor import DTensor, Replicate, Shard,
|
|
|
|
| 11 |
from torch.profiler import ProfilerActivity, profile
|
| 12 |
|
| 13 |
from .utils import ParallelDims, assert_params_equal, parallelize_llama4
|
|
|
|
| 7 |
import torch
|
| 8 |
import torch.distributed as dist
|
| 9 |
from optimizer.muon import Muon, get_default_muon_param_groups
|
| 10 |
+
from torch.distributed.tensor import (DTensor, Replicate, Shard,
|
| 11 |
+
distribute_tensor)
|
| 12 |
from torch.profiler import ProfilerActivity, profile
|
| 13 |
|
| 14 |
from .utils import ParallelDims, assert_params_equal, parallelize_llama4
|