Source code for lumin.nn.callbacks.pred_handlers

import numpy as np

import torch

from .callback import Callback


class PredHandler(Callback):
    r'''
    Default callback for predictions. Collects predictions over batches and returns them as stacked array
    '''

    def on_pred_begin(self) -> None:
        super().on_pred_begin()
        self.preds = []

    def on_pred_end(self) -> None: self.preds = torch.cat(self.preds)
    def get_preds(self) -> np.ndarray: return self.preds
    def on_forwards_end(self) -> None:
        if self.model.fit_params.state == 'test': self.preds.append(self.model.fit_params.y_pred.detach().cpu())