fix ci

Signed-off-by: YunLiu <[email protected]>

Author: KumoLiu

generation/image_to_image_translation/tutorial_segmentation_with_ddpm.ipynb

@@ -393,7 +393,7 @@
     "            seg = data_val[\"label\"].to(device)  # this is the ground truth segmentation\n",
     "            timesteps = torch.randint(0, 1000, (len(images),)).to(device)\n",
     "            with torch.no_grad():\n",
-    "                with autocast(enabled=True):\n",
+    "                with autocast(\"cuda\", enabled=True):\n",
     "                    noise = torch.randn_like(seg).to(device)\n",
     "                    noisy_seg = scheduler.add_noise(original_samples=seg, noise=noise, timesteps=timesteps)\n",
     "                    combined = torch.cat((images, noisy_seg), dim=1)\n",

modules/interpretability/cats_and_dogs.ipynb

@@ -317,7 +317,7 @@
     "    for data in dl:\n",
     "        inputs, labels = data[\"image\"].to(device), data[\"label\"].to(device)\n",
     "        optimizer.zero_grad()\n",
-    "        with torch.autocast(\"cuda\", enabled=use_amp):\n",
+    "        with torch.autocast(\"cuda\", enabled=use_amp) if use_amp else nullcontext():\n",
     "            outputs = model(inputs)\n",
     "            train_loss = criterion(outputs, labels)\n",
     "            acc += get_num_correct(outputs, labels)\n",
@@ -350,7 +350,7 @@
     "\n",
     "    for data in tqdm(dl):\n",
     "        images, labels = data[\"image\"].to(device), data[\"label\"].to(device)\n",
-    "        with torch.autocast(\"cuda\", enabled=use_amp):\n",
+    "        with torch.autocast(\"cuda\", enabled=use_amp) if use_amp else nullcontext():\n",
     "            outputs = model(images).detach()\n",
     "        y_pred = torch.cat([y_pred, outputs], dim=0)\n",
     "        y = torch.cat([y, labels], dim=0)\n",