def test_export_pytorch_model(self):
"""Tests loading of onnx model from a file"""
pytorch_model = PyTorchLinear()
dummy_input = torch.empty(10, 10)
with io.BytesIO() as f:
onnx_converter._export_pytorch_model(f, pytorch_model, dummy_input)
def test_from_onnx(self):
"""Tests construction of crypten model from onnx graph"""
pytorch_model = PyTorchLinear()
dummy_input = torch.empty(10, 10)
with io.BytesIO() as f:
f = onnx_converter._export_pytorch_model(f, pytorch_model, dummy_input)
f.seek(0)
crypten_model = onnx_converter.from_onnx(f)
self.assertTrue(hasattr(crypten_model, "encrypt"))
