def METHOD_NAME(self, runner):
    if runner.epoch == 0 and self.add_graph:
        if is_module_wrapper(runner.model):
            _model = runner.model.module
        else:
            _model = runner.model
        device = next(_model.parameters()).device
        data = next(iter(runner.data_loader))
        image = data[self.img_key][0:1].to(device)
        with torch.no_grad():
            self.writer.add_graph(_model, image)