diff --git a/notebooks/EfficientNet-example.ipynb b/notebooks/EfficientNet-example.ipynb index 7e7feb93a5..cfb8e79232 100644 --- a/notebooks/EfficientNet-example.ipynb +++ b/notebooks/EfficientNet-example.ipynb @@ -610,7 +610,7 @@ "source": [ "# The compiled module will have precision as specified by \"op_precision\".\n", "# Here, it will have FP16 precision.\n", - "trt_model_fp16 = torch_tensorrt.compile(model, inputs = [torch_tensorrt.Input((128, 3, 224, 224), dtype=torch.half)],\n", + "trt_model_fp16 = torch_tensorrt.compile(model.half(), inputs = [torch_tensorrt.Input((128, 3, 224, 224), dtype=torch.half)],\n", " enabled_precisions = {torch.half}, # Run with FP16\n", " workspace_size = 1 << 22\n", ")"