simple inference

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))