import torch

import torch.nn as nn

from tqdm import tqdm

from data import GIImageDataLoader

# TODO: other metrics

def evaluate(model: nn.Module, data_loader: GIImageDataLoader, criterion: nn.Module, device: str, use_fp16: bool) -> float:

    model.eval()

    model = model.to(device)

    criterion = criterion.to(device)

    total_loss = 0.0

    with torch.no_grad():

        for inputs, labels in tqdm(data_loader):

            inputs = inputs.to(device)

            labels = labels.to(device)

            with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=use_fp16):

                preds = model(inputs)

                loss = criterion(preds, labels)

            total_loss += loss.item() * inputs.size(0)

    return total_loss