def test_quantize_torchscript(tmpdir):
"""Test converting to torchscipt """
dm = RegressDataModule()
qmodel = RegressionModel()
qcb = QuantizationAwareTraining(input_compatible=False)
trainer = Trainer(callbacks=[qcb], default_root_dir=tmpdir, max_epochs=1)
trainer.fit(qmodel, datamodule=dm)
qmodel.to_torchscript()
def test_quantize_torchscript(tmpdir):
"""Test converting to torchscipt."""
dm = RegressDataModule()
qmodel = RegressionModel()
qcb = QuantizationAwareTraining(input_compatible=False)
trainer = Trainer(callbacks=[qcb], default_root_dir=tmpdir, max_epochs=1)
trainer.fit(qmodel, datamodule=dm)
batch = iter(dm.test_dataloader()).next()
qmodel(qmodel.quant(batch[0]))
tsmodel = qmodel.to_torchscript()
tsmodel(tsmodel.quant(batch[0]))