pytorch TorchScript

	model = BertClassifier()
	model.load_state_dict(torch.load('model.pth', map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu')))
	model.eval()	
	tokenizer = BertTokenizer.from_pretrained('hfl/chinese-roberta-wwm-ext')
	
    token = tokenizer(text, add_special_tokens=True, padding='max_length', truncation=True, max_length=128)
    input_ids = token['input_ids']
    token_type_ids = token['token_type_ids']
    attention_mask = token['attention_mask']

    input_ids = torch.tensor([input_ids], dtype=torch.long)
    attention_mask = torch.tensor([attention_mask], dtype=torch.long)
    token_type_ids = torch.tensor([token_type_ids], dtype=torch.long)

    jit_sample = (input_ids, attention_mask, token_type_ids)
    module = torch.jit.trace(model, jit_sample)
    module.save('pthmodes/model_jit.pt')
    jit_predict = module(
        input_ids,
        attention_mask,
        token_type_ids,
    )   

    model_predict = model(
        input_ids,
        attention_mask,
        token_type_ids,
    )