gc

Contents

`gc`#

class fkat.pytorch.callbacks.gc.ManualGc(schedule: Optional[Schedule] = None, stats: Optional[dict[str, list[str]]] = None)[source]#

maybe_collect() → None[source]#

maybe_gc(trainer: Trainer, stage: str, batch_idx: int) → None[source]#

Perform garbage collection if conditions are met.

Parameters:: batch_idx (int) – Current batch index

on_predict_batch_end(trainer: Trainer, pl_module: LightningModule, outputs: Any, batch: Any, batch_idx: int, dataloader_idx: int = 0) → None[source]#: Perform GC after prediction batch if needed.

on_predict_epoch_end(trainer: Trainer, pl_module: LightningModule) → None[source]#: Perform GC after predict epoch.

on_test_batch_end(trainer: Trainer, pl_module: LightningModule, outputs: Optional[Union[Tensor, Mapping[str, Any]]], batch: Any, batch_idx: int, dataloader_idx: int = 0) → None[source]#: Perform GC after test batch if needed.

on_test_epoch_end(trainer: Trainer, pl_module: LightningModule) → None[source]#: Perform GC after test epoch.

on_train_batch_end(trainer: Trainer, pl_module: LightningModule, outputs: Optional[Union[Tensor, Mapping[str, Any]]], batch: Any, batch_idx: int) → None[source]#: Perform GC after training batch if needed.

on_train_epoch_end(trainer: Trainer, pl_module: LightningModule) → None[source]#: Perform GC after training epoch.

on_validation_batch_end(trainer: Trainer, pl_module: LightningModule, outputs: Optional[Union[Tensor, Mapping[str, Any]]], batch: Any, batch_idx: int, dataloader_idx: int = 0) → None[source]#: Perform GC after validation batch if needed.

on_validation_epoch_end(trainer: Trainer, pl_module: LightningModule) → None[source]#: Perform GC after validation epoch.

setup(trainer: Trainer, pl_module: LightningModule, stage: str) → None[source]#: Called when fit, validate, test, predict, or tune begins.