Skip to content

Commit d0ad145

Browse files
authored
Update __init__.py
1 parent 3d3323d commit d0ad145

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

src/lightning/pytorch/callbacks/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@
3232
from lightning.pytorch.callbacks.stochastic_weight_avg import StochasticWeightAveraging
3333
from lightning.pytorch.callbacks.throughput_monitor import ThroughputMonitor
3434
from lightning.pytorch.callbacks.timer import Timer
35-
from lightning.pytorch.callbacks.weight_averaging import WeightAveraging
35+
from lightning.pytorch.callbacks.weight_averaging import EMAWeightAveraging, WeightAveraging
3636

3737
__all__ = [
3838
"BackboneFinetuning",
@@ -59,5 +59,6 @@
5959
"ThroughputMonitor",
6060
"Timer",
6161
"TQDMProgressBar",
62+
"EMAWeightAveraging",
6263
"WeightAveraging",
6364
]

0 commit comments

Comments
 (0)