simple inference

inference in pytorch

November 12, 2025 · 1 min read · tech, programming, machine-learning

import torch
import torch.nn as nn

# input, model, and output are on the GPU
def solve(input: torch.Tensor, model: nn.Module, output: torch.Tensor):
    model.eval()
    with torch.no_grad():
        output.copy_(model(input))