Kernels
wyldecat Claude Opus 4.6 commited on
Commit
96b287c
·
1 Parent(s): bdada12

Apply pre-commit formatting (isort) [skip-build]

Browse files

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Files changed (1) hide show
  1. 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, distribute_tensor
 
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