GPUAccelerator

class pytorch_lightning.accelerators.GPUAccelerator(precision_plugin, training_type_plugin)

Bases: pytorch_lightning.accelerators.accelerator.Accelerator

Accelerator for GPU devices.

Parameters

• precision_plugin (PrecisionPlugin) – the plugin to handle precision-specific parts

• training_type_plugin (TrainingTypePlugin) – the plugin to handle different training routines

on_train_start()

Hook to do something upon the training start

Return type

None

setup(trainer, model)

Raises

MisconfigurationException – If the selected device is not GPU.

Return type

None

teardown()

This method is called to teardown the training process. It is the right place to release memory and free other ressources.

By default we add a barrier here to synchronize processes before returning control back to the caller.

Return type

None

to_device(step_kwargs)

Pushes the batch to the root device

Return type

Dict[str, Union[Any, int]]