From db379e991da5c9e99fe38500537cf8e86beac02e Mon Sep 17 00:00:00 2001
From: Valerio Sofi <v.sofi@nebuly.ai>
Date: Mon, 12 Sep 2022 12:18:20 +0200
Subject: [PATCH] fix type when tensorflow is not installed (#97)

* fix type when tensorflow is not installed

* fix compressor bug

* fix torch tensorrt

Co-authored-by: Valerio Sofi <v.sofi@nebuly.ai>
---
 nebullvm/compressors/scripts/__init__.py |  0
 nebullvm/optimizers/tensor_rt.py         | 10 ++++++++--
 nebullvm/utils/optional_modules.py       |  8 +++++---
 3 files changed, 13 insertions(+), 5 deletions(-)
 create mode 100644 nebullvm/compressors/scripts/__init__.py

diff --git a/nebullvm/compressors/scripts/__init__.py b/nebullvm/compressors/scripts/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/nebullvm/optimizers/tensor_rt.py b/nebullvm/optimizers/tensor_rt.py
index e2f54369..fbc3eb06 100644
--- a/nebullvm/optimizers/tensor_rt.py
+++ b/nebullvm/optimizers/tensor_rt.py
@@ -335,7 +335,7 @@ def optimize_from_torch(
                 dataset,
                 batch_size=dataset.batch_size,
                 shuffle=False,
-                num_workers=1,
+                num_workers=0,
             )
 
             calibrator = torch_tensorrt.ptq.DataLoaderCalibrator(
@@ -350,8 +350,14 @@ def optimize_from_torch(
         ):
             return None  # Dynamic quantization is not supported on tensorRT
 
+        try:
+            torch.jit.script(torch_model.eval())
+            model = torch_model
+        except Exception:
+            model = torch.jit.trace(torch_model, input_data.get_list(1)[0])
+
         trt_model = torch_tensorrt.compile(
-            torch_model.eval(),
+            model.eval(),
             inputs=[
                 torch_tensorrt.Input(
                     (model_params.batch_size, *input_info.size),
diff --git a/nebullvm/utils/optional_modules.py b/nebullvm/utils/optional_modules.py
index a57f8f5f..e8c37b84 100644
--- a/nebullvm/utils/optional_modules.py
+++ b/nebullvm/utils/optional_modules.py
@@ -3,14 +3,16 @@
 from nebullvm.installers.installers import install_tf2onnx, install_tensorflow
 from nebullvm.utils.general import check_module_version
 
+NoneType = type(None)
+
 
 class Keras:
-    Model = None
+    Model = NoneType
 
 
 class Tensorflow:
-    Module = None
-    Tensor = None
+    Module = NoneType
+    Tensor = NoneType
     keras = Keras()
 
     @staticmethod