一个小错误

pull/50/head
forXuyx 2025-03-22 23:31:31 +08:00
parent 4a80ca7c78
commit fc64a31c0a
1 changed files with 1 additions and 1 deletions

View File

@ -86,7 +86,7 @@ def train_epoch(epoch, wandb):
if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0): if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
model.eval() model.eval()
moe_path = '_moe' if model_config.use_moe else '' moe_path = '_moe' if model_config.use_moe else ''
ckp = f'{args.save_dir}/pretrain_vlm_{model_config.dim}{moe_path}.pth' ckp = f'{args.save_dir}/pretrain_videolm_{model_config.dim}{moe_path}.pth'
if isinstance(model, torch.nn.parallel.DistributedDataParallel): if isinstance(model, torch.nn.parallel.DistributedDataParallel):
state_dict = model.module.state_dict() state_dict = model.module.state_dict()
else: else: