Simple Inference

1
import torch
2
import torch.nn as nn
3

4

5
# input, model, and output are on the GPU
6
def solve(input: torch.Tensor, model: nn.Module, output: torch.Tensor):
7
    weights = model.weight.T
8
    res = input @ weights
9
    if model.bias is not None:
10
        res = res + model.bias
11
    output.copy_(res)