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

@@ -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",