From 952d778dd897a4888460b185ce65d343fbf9b037 Mon Sep 17 00:00:00 2001
From: Noah Luna <15202580+ngrayluna@users.noreply.github.com>
Date: Tue, 20 Jun 2023 09:27:52 -0700
Subject: [PATCH] Changed RayTune import statement. (#427)

* Changed RayTune import statement.

* fix CB integration

---------

Co-authored-by: Thomas Capelle <tcapelle@pm.me>
---
 colabs/raytune/RayTune_with_wandb.ipynb | 17 +++--------------
 1 file changed, 3 insertions(+), 14 deletions(-)

diff --git a/colabs/raytune/RayTune_with_wandb.ipynb b/colabs/raytune/RayTune_with_wandb.ipynb
index c70dac71..7d2ff47b 100644
--- a/colabs/raytune/RayTune_with_wandb.ipynb
+++ b/colabs/raytune/RayTune_with_wandb.ipynb
@@ -1,7 +1,6 @@
 {
  "cells": [
   {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -124,8 +123,7 @@
     "import random\n",
     "import numpy as np\n",
     "from ray import tune\n",
-    "from ray.tune.logger import DEFAULT_LOGGERS\n",
-    "from ray.air.callbacks.wandb import WandbLoggerCallback\n",
+    "from ray.air.integrations.wandb import WandbLoggerCallback, wandb_setup\n",
     "import torch\n",
     "import torch.optim as optim\n",
     "import wandb"
@@ -207,9 +205,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "@wandb_mixin\n",
     "def train_mnist(config):\n",
-    "\n",
+    "    # Setup wandb\n",
+    "    wandb = wandb_setup(config)\n",
     "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
     "    train_loader, test_loader = get_data_loaders()\n",
     "\n",
@@ -262,15 +260,6 @@
     "[the documentation for the `WandbLoggerCallback`](https://docs.ray.io/en/master/tune/tutorials/tune-wandb.html)."
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "wandb.login()"
-   ]
-  },
   {
    "attachments": {},
    "cell_type": "markdown",