Fix error on TensorRT export with float workspace value (ultralytic…

…s#17352)
triple-Mu · Nov 5, 2024 · d0abd95 · d0abd95
1 parent 603fa84
commit d0abd95
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py
@@ -791,7 +791,7 @@ def export_engine(self, dla=None, prefix=colorstr("TensorRT:")):
                 LOGGER.warning(f"{prefix} WARNING ⚠️ 'dynamic=True' model requires max batch size, i.e. 'batch=16'")
             profile = builder.create_optimization_profile()
             min_shape = (1, shape[1], 32, 32)  # minimum input shape
-            max_shape = (*shape[:2], *(max(1, self.args.workspace) * d for d in shape[2:]))  # max input shape
+            max_shape = (*shape[:2], *(int(max(1, self.args.workspace) * d) for d in shape[2:]))  # max input shape
             for inp in inputs:
                 profile.set_shape(inp.name, min=min_shape, opt=shape, max=max_shape)
             config.add_optimization_profile(profile)