From e57809c2d9de150c7486054987dde36b592a3c21 Mon Sep 17 00:00:00 2001
From: kcelia <celia.kherfallah@zama.ai>
Date: Mon, 23 Sep 2024 16:18:08 +0200
Subject: [PATCH] chore: update cifar finetuning

---
 .../cifar/cifar_brevitas_finetuning/CifarInFhe.ipynb        | 6 ++++--
 .../CifarInFheWithSmallerAccumulators.ipynb                 | 6 ++++--
 .../CifarQuantizationAwareTraining.ipynb                    | 6 ++++--
 .../cifar_brevitas_finetuning/FromImageNetToCifar.ipynb     | 6 ++++--
 .../cifar/cifar_brevitas_finetuning/cifar_utils.py          | 2 +-
 5 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/use_case_examples/cifar/cifar_brevitas_finetuning/CifarInFhe.ipynb b/use_case_examples/cifar/cifar_brevitas_finetuning/CifarInFhe.ipynb
index da3aa471af..099355d03b 100644
--- a/use_case_examples/cifar/cifar_brevitas_finetuning/CifarInFhe.ipynb
+++ b/use_case_examples/cifar/cifar_brevitas_finetuning/CifarInFhe.ipynb
@@ -63,8 +63,10 @@
     "bit = 5\n",
     "seed = 42\n",
     "\n",
-    "use_gpu_if_available = False\n",
-    "device = \"cuda\" if use_gpu_if_available and check_gpu_available() else \"cpu\"\n",
+    "# To use the CUDA backend for the Concrete compiler, install the GPU-supported version.\n",
+    "# Refer to the official guide:\n",
+    "# https://github.com/zama-ai/concrete-ml/blob/main/docs/guides/using_gpu.md\n",
+    "device = \"cuda\" if check_gpu_available() else \"cpu\"\n",
     "\n",
     "print(f\"Device Type: {device}\")"
    ]
diff --git a/use_case_examples/cifar/cifar_brevitas_finetuning/CifarInFheWithSmallerAccumulators.ipynb b/use_case_examples/cifar/cifar_brevitas_finetuning/CifarInFheWithSmallerAccumulators.ipynb
index b76edd99e8..ffe2834db6 100644
--- a/use_case_examples/cifar/cifar_brevitas_finetuning/CifarInFheWithSmallerAccumulators.ipynb
+++ b/use_case_examples/cifar/cifar_brevitas_finetuning/CifarInFheWithSmallerAccumulators.ipynb
@@ -61,8 +61,10 @@
     "seed = 42\n",
     "rounding_thresholds_bits = [8, 7, 6, 5, 3]\n",
     "\n",
-    "use_gpu_if_available = False\n",
-    "device = \"cuda\" if use_gpu_if_available and check_gpu_available() else \"cpu\""
+    "# To use the CUDA backend for the Concrete compiler, install the GPU-supported version.\n",
+    "# Refer to the official guide:\n",
+    "# https://github.com/zama-ai/concrete-ml/blob/main/docs/guides/using_gpu.md\n",
+    "device = \"cuda\" if check_gpu_available() else \"cpu\""
    ]
   },
   {
diff --git a/use_case_examples/cifar/cifar_brevitas_finetuning/CifarQuantizationAwareTraining.ipynb b/use_case_examples/cifar/cifar_brevitas_finetuning/CifarQuantizationAwareTraining.ipynb
index 6e834d91f1..a782deea86 100644
--- a/use_case_examples/cifar/cifar_brevitas_finetuning/CifarQuantizationAwareTraining.ipynb
+++ b/use_case_examples/cifar/cifar_brevitas_finetuning/CifarQuantizationAwareTraining.ipynb
@@ -89,8 +89,10 @@
    "source": [
     "bit = 5\n",
     "\n",
-    "use_gpu_if_available = False\n",
-    "device = \"cuda\" if use_gpu_if_available and check_gpu_available() else \"cpu\"\n",
+    "# To use the CUDA backend for the Concrete compiler, install the GPU-supported version.\n",
+    "# Refer to the official guide:\n",
+    "# https://github.com/zama-ai/concrete-ml/blob/main/docs/guides/using_gpu.md\n",
+    "device = \"cuda\" if check_gpu_available() else \"cpu\"\n",
     "\n",
     "print(f\"Device Type: {device}\")"
    ]
diff --git a/use_case_examples/cifar/cifar_brevitas_finetuning/FromImageNetToCifar.ipynb b/use_case_examples/cifar/cifar_brevitas_finetuning/FromImageNetToCifar.ipynb
index ed89359f25..e272f509b9 100644
--- a/use_case_examples/cifar/cifar_brevitas_finetuning/FromImageNetToCifar.ipynb
+++ b/use_case_examples/cifar/cifar_brevitas_finetuning/FromImageNetToCifar.ipynb
@@ -65,8 +65,10 @@
    "source": [
     "dataset_name = \"CIFAR_100\"\n",
     "\n",
-    "use_gpu_if_available = False\n",
-    "device = \"cuda\" if use_gpu_if_available and check_gpu_available() else \"cpu\"\n",
+    "# To use the CUDA backend for the Concrete compiler, install the GPU-supported version.\n",
+    "# Refer to the official guide:\n",
+    "# https://github.com/zama-ai/concrete-ml/blob/main/docs/guides/using_gpu.md#usage-preqreuisites\n",
+    "device = \"cuda\" if check_gpu_available() else \"cpu\"\n",
     "\n",
     "param_c10 = {\n",
     "    \"output_size\": 10,\n",
diff --git a/use_case_examples/cifar/cifar_brevitas_finetuning/cifar_utils.py b/use_case_examples/cifar/cifar_brevitas_finetuning/cifar_utils.py
index ac335af7d7..be40ba9eae 100644
--- a/use_case_examples/cifar/cifar_brevitas_finetuning/cifar_utils.py
+++ b/use_case_examples/cifar/cifar_brevitas_finetuning/cifar_utils.py
@@ -457,7 +457,7 @@ def fhe_compatibility(model: Callable, data: DataLoader, device: str) -> Callabl
         torch_inputset=data,
         show_mlir=False,
         output_onnx_file="test.onnx",
-        configuration=Configuration(use_gpu=(device == "cuda")),
+        device=device,
     )
 
     return qmodel