34 lines
796 B
Python
34 lines
796 B
Python
import torch
|
|
from typing import Any
|
|
from pytorch_lightning.utilities.apply_func import move_data_to_device
|
|
|
|
|
|
class Accelerator(object):
|
|
|
|
def __init__(self, trainer):
|
|
self.trainer = trainer
|
|
|
|
def setup(self, model):
|
|
pass
|
|
|
|
def teardown(self):
|
|
pass
|
|
|
|
def batch_to_device(self, batch: Any, device: torch.device):
|
|
model = self.trainer.get_model()
|
|
if model is not None:
|
|
return model.transfer_batch_to_device(batch, device)
|
|
return move_data_to_device(batch, device)
|
|
|
|
def training_step_end(self, output):
|
|
return output
|
|
|
|
def test_step_end(self, output):
|
|
return output
|
|
|
|
def validation_step_end(self, output):
|
|
return output
|
|
|
|
def process_dataloader(self, dataloader):
|
|
return dataloader
|