small changes

grazder · Mar 14, 2024 · 04d035c · 04d035c
1 parent dfa1227
commit 04d035c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/torchDF/model_onnx_export.py b/torchDF/model_onnx_export.py
@@ -167,10 +167,13 @@ def custom_identity(g: jit_utils.GraphContext, X):
 
 
 def main(args):
+    torch.set_num_threads(1)
+    torch.set_num_interop_threads(1)
+
     if args.minimal:
         streaming_pipeline = TorchDFMinimalPipeline(device="cpu")
     else:
-        streaming_pipeline = TorchDFPipeline(device="cpu")
+        streaming_pipeline = TorchDFPipeline(device="cpu", always_apply_all_stages=True)
 
     frame_size = streaming_pipeline.hop_size
     input_names = streaming_pipeline.input_names
@@ -212,7 +215,7 @@ def main(args):
     input_shapes_dict = {x: y.shape for x, y in input_features_onnx.items()}
 
     # Simplify not working for not minimal!
-    if args.simplify and args.minimal:
+    if args.simplify:
         # raise NotImplementedError("Simplify not working for flatten states!")
         onnx_simplify(args.output_path, input_features_onnx, input_shapes_dict)
         logger.info(f"Model simplified! {args.output_path}")

diff --git a/torchDF/torch_df_streaming_minimal.py b/torchDF/torch_df_streaming_minimal.py
@@ -16,14 +16,12 @@
 from df import init_df
 
 from functools import partial
-from typing import Callable, Iterable, List, Optional, Tuple, Union
+from typing import Callable, List, Optional, Tuple
 
 import torch
 from torch import Tensor, nn
 import torch.nn.functional as F
 
-from loguru import logger
-
 from df.config import Csv, DfParams, config
 from df.modules import Conv2dNormAct, ConvTranspose2dNormAct