Function test_model_inference

SwissArmyTransformer/tests/test_inference.py:7–11

()

def test_model_inference():
    from sat.model import RobertaModel, AutoModel
    model, args1 = RobertaModel.from_pretrained('roberta-base')
    x = torch.tensor([[1,2,3]], device='cuda')
    a = model(input_ids=x, position_ids=x, attention_mask=None)

def test_model_inference_create():
    from sat.model import RobertaModel, AutoModel