bert_torch.py

import torch
from shark.shark_runner import SharkInference

from bert_pytorch import BERT

torch.manual_seed(0)


class BERT_torch(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = BERT(512)

    def forward(self, tokens):
        return self.model.forward(tokens, tokens)



test_input = torch.randint(2,(1,128)).to(torch.int32)

shark_module = SharkInference(
    BERT_torch(), (test_input,), device="cpu", jit_trace=True
)

results = shark_module.forward((test_input,))