5050import lightning .pytorch as pl
5151import torch
5252from lightning .pytorch import Callback
53- from lightning .pytorch .trainer .states import TrainerFn
5453from lightning .pytorch .utilities .types import STEP_OUTPUT
5554from torch import nn
5655from torchvision .transforms .v2 import Compose , Normalize , Resize
@@ -143,7 +142,6 @@ def __init__(
143142 self .visualizer = self ._resolve_visualizer (visualizer )
144143
145144 self ._input_size : tuple [int , int ] | None = None
146- self ._is_setup = False
147145
148146 @property
149147 def name (self ) -> str :
@@ -154,33 +152,6 @@ def name(self) -> str:
154152 """
155153 return self .__class__ .__name__
156154
157- def setup (self , stage : str | None = None ) -> None :
158- """Set up the model if not already done.
159-
160- This method ensures the model is built by calling ``_setup()`` if needed.
161-
162- Args:
163- stage (str | None, optional): Current stage of training.
164- Defaults to ``None``.
165- """
166- if getattr (self , "model" , None ) is None or not self ._is_setup :
167- self ._setup ()
168- if isinstance (stage , TrainerFn ):
169- # only set the flag if the stage is a TrainerFn, which means the
170- # setup has been called from a trainer
171- self ._is_setup = True
172-
173- def _setup (self ) -> None :
174- """Set up the model architecture.
175-
176- This method should be overridden by subclasses to build their model
177- architecture. It is called by ``setup()`` when the model needs to be
178- initialized.
179-
180- This is useful when the model cannot be fully initialized in ``__init__``
181- because it requires data-dependent parameters.
182- """
183-
184155 def configure_callbacks (self ) -> Sequence [Callback ] | Callback :
185156 """Configure callbacks for the model.
186157
0 commit comments