DoublePrecisionPlugin

class pytorch_lightning.plugins.precision.DoublePrecisionPlugin

Bases: pytorch_lightning.plugins.precision.precision_plugin.PrecisionPlugin

Plugin for training with double (torch.float64) precision.

connect(model, optimizers, lr_schedulers)

Converts the model to double precision and wraps the training_step, validation_step, test_step, predict_step, and forward methods to convert incoming floating point data to double. Does not alter optimizers or lr_schedulers.

Return type

Tuple[Module, List[Optimizer], List[Any]]

post_dispatch()

Hook to do something after the training/evaluation/prediction finishes.

Return type

None

predict_step_context()

A context manager to change the default tensor type. See: torch.set_default_tensor_type()

Return type

Generator[None, None, None]

test_step_context()

A context manager to change the default tensor type. See: torch.set_default_tensor_type()

Return type

Generator[None, None, None]

train_step_context()

A context manager to change the default tensor type. See: torch.set_default_tensor_type()

Return type

Generator[None, None, None]

val_step_context()

A context manager to change the default tensor type. See: torch.set_default_tensor_type()

Return type

Generator[None, None, None]