This commit is contained in:
Srini Iyer 2025-03-11 23:53:01 +00:00
parent 6160e83d41
commit a3973339f2

View file

@ -4,7 +4,6 @@ import json
import logging import logging
import os import os
import re import re
import typer
import fsspec import fsspec
import s3fs import s3fs
@ -13,6 +12,7 @@ import torch.distributed as dist
import torch.distributed.checkpoint as dcp import torch.distributed.checkpoint as dcp
import torch.nn as nn import torch.nn as nn
import torch.optim.optimizer import torch.optim.optimizer
import typer
from pydantic import BaseModel, ConfigDict from pydantic import BaseModel, ConfigDict
from torch.distributed._tensor import DeviceMesh from torch.distributed._tensor import DeviceMesh
from torch.distributed.checkpoint.format_utils import dcp_to_torch_save from torch.distributed.checkpoint.format_utils import dcp_to_torch_save