Gemfury - distributed/fsdp/_fully_shard/__init__.py · team-dmm / torch v2.7.1 - python package

team-dmm / torch python

Repository URL to install this package:

Details

torch / distributed / fsdp / _fully_shard / __init__.py

from ._fsdp_api import CPUOffloadPolicy, MixedPrecisionPolicy, OffloadPolicy
from ._fully_shard import (
    FSDPModule,
    fully_shard,
    register_fsdp_forward_method,
    UnshardHandle,
)


__all__ = [
    "CPUOffloadPolicy",
    "FSDPModule",
    "fully_shard",
    "MixedPrecisionPolicy",
    "OffloadPolicy",
    "register_fsdp_forward_method",
    "UnshardHandle",
]