Skip to content

Commit 378d88c

Browse files
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
1 parent 40b7283 commit 378d88c

File tree

1 file changed

+7
-2
lines changed
  • src/lightning/pytorch/strategies

1 file changed

+7
-2
lines changed

src/lightning/pytorch/strategies/fsdp.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
import logging
1515
from contextlib import contextmanager, nullcontext
1616
from datetime import timedelta
17-
from typing import Any, Callable, Dict, Generator, List, Mapping, Optional, Type, TYPE_CHECKING, Union, Literal
17+
from typing import Any, Callable, Dict, Generator, List, Literal, Mapping, Optional, Type, TYPE_CHECKING, Union
1818

1919
import torch
2020
from torch import Tensor
@@ -60,7 +60,12 @@
6060
from lightning.pytorch.utilities.rank_zero import rank_zero_info, rank_zero_only
6161

6262
if TYPE_CHECKING:
63-
from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload, FullyShardedDataParallel, MixedPrecision, ShardingStrategy
63+
from torch.distributed.fsdp.fully_sharded_data_parallel import (
64+
CPUOffload,
65+
FullyShardedDataParallel,
66+
MixedPrecision,
67+
ShardingStrategy,
68+
)
6469

6570
if _TORCH_GREATER_EQUAL_2_0:
6671
from torch.distributed.fsdp.wrap import _FSDPPolicy

0 commit comments

Comments
 (0)