diff --git a/colabs/pytorch-lightning/Profile_PyTorch_Code.ipynb b/colabs/pytorch-lightning/Profile_PyTorch_Code.ipynb
index 573e035a..79815c01 100644
--- a/colabs/pytorch-lightning/Profile_PyTorch_Code.ipynb
+++ b/colabs/pytorch-lightning/Profile_PyTorch_Code.ipynb
@@ -350,7 +350,7 @@
     "    with profiler:\n",
     "        profiler_callback = TorchTensorboardProfilerCallback(profiler)\n",
     "\n",
-    "        trainer = pl.Trainer(gpus=1, max_epochs=1, max_steps=total_steps,\n",
+    "        trainer = pl.Trainer(max_epochs=1, max_steps=total_steps,\n",
     "                            logger=pl.loggers.WandbLogger(log_model=True, save_code=True),\n",
     "                            callbacks=[profiler_callback], precision=wandb.config.precision)\n",
     "\n",