diff options
| author | One <imone@tuta.io> | 2025-07-21 18:40:40 +0800 |
|---|---|---|
| committer | One <imone@tuta.io> | 2025-07-21 18:40:40 +0800 |
| commit | 171e2fcde636bcb7e6c0073a9983ed5252f04753 (patch) | |
| tree | d7844d28ad5f289c25a046e58ec9d20216cfba44 /pretrain.py | |
| parent | bd6222774edcec1608a6842d0b06a637a4acef59 (diff) | |
Update
Diffstat (limited to 'pretrain.py')
| -rw-r--r-- | pretrain.py | 3 |
1 files changed, 1 insertions, 2 deletions
diff --git a/pretrain.py b/pretrain.py index b939318..245cb5c 100644 --- a/pretrain.py +++ b/pretrain.py @@ -16,7 +16,6 @@ import coolname import hydra import pydantic from omegaconf import DictConfig -from wandb.util import make_artifact_name_safe from adam_atan2 import AdamATan2 from puzzle_dataset import PuzzleDataset, PuzzleDatasetConfig, PuzzleDatasetMetadata @@ -126,7 +125,7 @@ def create_model(config: PretrainConfig, train_metadata: PuzzleDatasetMetadata, model: nn.Module = model_cls(model_cfg) model = loss_head_cls(model, **config.arch.loss.__pydantic_extra__) # type: ignore if "DISABLE_COMPILE" not in os.environ: - model = torch.compile(model, dynamic=False, fullgraph=True) # type: ignore + model = torch.compile(model, dynamic=False) # type: ignore # Broadcast parameters from rank 0 if world_size > 1: |