Skip to content

Commit

Permalink
update
Browse files Browse the repository at this point in the history
  • Loading branch information
rlsu9 committed Jan 5, 2025
1 parent 0b1a50d commit 14afa22
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions tests/test_save_checkpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,6 @@
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

from fastvideo.utils.checkpoint import save_checkpoint
from fastvideo.utils.fsdp_util import get_dit_fsdp_kwargs


@pytest.fixture(scope="module", autouse=True)
def setup_distributed():
Expand All @@ -28,6 +25,9 @@ def setup_distributed():
def test_save_and_remove_checkpoint():
from fastvideo.models.mochi_hf.modeling_mochi import \
MochiTransformer3DModel
from fastvideo.utils.checkpoint import save_checkpoint
from fastvideo.utils.fsdp_util import get_dit_fsdp_kwargs

transformer = MochiTransformer3DModel(num_layers=0)
fsdp_kwargs, _ = get_dit_fsdp_kwargs(transformer, "none")
transformer = FSDP(transformer, **fsdp_kwargs)
Expand Down

0 comments on commit 14afa22

Please sign in to comment.