test_linalg.py

import torch
import torch_mlir


class ANET(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.alexnet = torch.hub.load('pytorch/vision:v0.10.0', 'alexnet', pretrained=True)
        self.alexnet.eval()

    def forward(self, image):
        return self.alexnet(image)


alexnet_model = ANET()
# batch, channels, height, width
rnd_inp = torch.randn(1,3,224,224)
module = torch_mlir.compile(alexnet_model, (rnd_inp,), output_type="linalg-on-tensors")
module.dump()