We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Fastai WandbCallback with multilabel datasets
* Weights and Biases version: 0.8.9
* Python version: 3.6.9
* Operating System: Linux-4.14.133-88.112.amzn1.x86_64-x86_64-with-Ubuntu-16.04-xenial

### Description

I'm trying to see prediction samples for fastai image classification, where the images are multilabel. I'd like  it to be ran on the validation data at the end of each epoch and to see captions of ground truth/predictions.

When I use a multilabel dataset, instead of ground truth/prediction captions for each image, I get a three copies of each image side by side, where each is labeled, "Input data", "Prediction", "Ground Truth" (but not the actual categories)

### What I Did

`callback_fns += [partial(WandbCallback, input_type='images', monitor='acc_thresholded')]`



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `wandb/fastai/__init__.py`
Content:
```
1 '''
2 This module hooks fast.ai Learners to Weights & Biases through a callback.
3 Requested logged data can be configured through the callback constructor.
4 
5 Examples:
6     WandbCallback can be used when initializing the Learner::
7 
8     ```
9         from wandb.fastai import WandbCallback
10         [...]
11         learn = Learner(data, ..., callback_fns=WandbCallback)
12         learn.fit(epochs)
13     ```
14 
15     Custom parameters can be given using functools.partial::
16 
17     ```
18         from wandb.fastai import WandbCallback
19         from functools import partial
20         [...]
21         learn = Learner(data, ..., callback_fns=partial(WandbCallback, ...))
22         learn.fit(epochs)
23     ```
24 
25     Finally, it is possible to use WandbCallback only when starting
26     training. In this case it must be instantiated::
27 
28     ```
29         learn.fit(..., callbacks=WandbCallback(learn))
30     ```
31 
32     or, with custom parameters::
33 
34     ```
35         learn.fit(..., callbacks=WandbCallback(learn, ...))
36     ```
37 '''
38 import wandb
39 import fastai
40 from fastai.callbacks import TrackerCallback
41 from pathlib import Path
42 import random
43 try:
44     import matplotlib
45     matplotlib.use('Agg')  # non-interactive backend (avoid tkinter issues)
46     import matplotlib.pyplot as plt
47 except:
48     print('Warning: matplotlib required if logging sample image predictions')
49 
50 
51 class WandbCallback(TrackerCallback):
52     """
53     Automatically saves model topology, losses & metrics.
54     Optionally logs weights, gradients, sample predictions and best trained model.
55 
56     Args:
57         learn (fastai.basic_train.Learner): the fast.ai learner to hook.
58         log (str): "gradients", "parameters", "all", or None. Losses & metrics are always logged.
59         save_model (bool): save model at the end of each epoch. It will also load best model at the end of training.
60         monitor (str): metric to monitor for saving best model. None uses default TrackerCallback monitor value.
61         mode (str): "auto", "min" or "max" to compare "monitor" values and define best model.
62         input_type (str): "images" or None. Used to display sample predictions.
63         validation_data (list): data used for sample predictions if input_type is set.
64         predictions (int): number of predictions to make if input_type is set and validation_data is None.
65         seed (int): initialize random generator for sample predictions if input_type is set and validation_data is None.
66     """
67 
68     # Record if watch has been called previously (even in another instance)
69     _watch_called = False
70 
71     def __init__(self,
72                  learn,
73                  log="gradients",
74                  save_model=True,
75                  monitor=None,
76                  mode='auto',
77                  input_type=None,
78                  validation_data=None,
79                  predictions=36,
80                  seed=12345):
81 
82         # Check if wandb.init has been called
83         if wandb.run is None:
84             raise ValueError(
85                 'You must call wandb.init() before WandbCallback()')
86 
87         # Adapted from fast.ai "SaveModelCallback"
88         if monitor is None:
89             # use default TrackerCallback monitor value
90             super().__init__(learn, mode=mode)
91         else:
92             super().__init__(learn, monitor=monitor, mode=mode)
93         self.save_model = save_model
94         self.model_path = Path(wandb.run.dir) / 'bestmodel.pth'
95 
96         self.log = log
97         self.input_type = input_type
98         self.best = None
99 
100         # Select items for sample predictions to see evolution along training
101         self.validation_data = validation_data
102         if input_type and not self.validation_data:
103             wandbRandom = random.Random(seed)  # For repeatability
104             predictions = min(predictions, len(learn.data.valid_ds))
105             indices = wandbRandom.sample(range(len(learn.data.valid_ds)),
106                                          predictions)
107             self.validation_data = [learn.data.valid_ds[i] for i in indices]
108 
109     def on_train_begin(self, **kwargs):
110         "Call watch method to log model topology, gradients & weights"
111 
112         # Set self.best, method inherited from "TrackerCallback" by "SaveModelCallback"
113         super().on_train_begin()
114 
115         # Ensure we don't call "watch" multiple times
116         if not WandbCallback._watch_called:
117             WandbCallback._watch_called = True
118 
119             # Logs model topology and optionally gradients and weights
120             wandb.watch(self.learn.model, log=self.log)
121 
122     def on_epoch_end(self, epoch, smooth_loss, last_metrics, **kwargs):
123         "Logs training loss, validation loss and custom metrics & log prediction samples & save model"
124 
125         if self.save_model:
126             # Adapted from fast.ai "SaveModelCallback"
127             current = self.get_monitor_value()
128             if current is not None and self.operator(current, self.best):
129                 print(
130                     'Better model found at epoch {} with {} value: {}.'.format(
131                         epoch, self.monitor, current))
132                 self.best = current
133 
134                 # Save within wandb folder
135                 with self.model_path.open('wb') as model_file:
136                     self.learn.save(model_file)
137 
138         # Log sample predictions if learn.predict is available
139         if self.validation_data:
140             try:
141                 self._wandb_log_predictions()
142             except FastaiError as e:
143                 wandb.termwarn(e.message)
144                 self.validation_data = None  # prevent from trying again on next loop
145             except Exception as e:
146                 wandb.termwarn("Unable to log prediction samples.\n{}".format(e))
147                 self.validation_data=None  # prevent from trying again on next loop
148 
149         # Log losses & metrics
150         # Adapted from fast.ai "CSVLogger"
151         logs = {
152             name: stat
153             for name, stat in list(
154                 zip(self.learn.recorder.names, [epoch, smooth_loss] +
155                     last_metrics))
156         }
157         wandb.log(logs)
158 
159     def on_train_end(self, **kwargs):
160         "Load the best model."
161 
162         if self.save_model:
163             # Adapted from fast.ai "SaveModelCallback"
164             if self.model_path.is_file():
165                 with self.model_path.open('rb') as model_file:
166                     self.learn.load(model_file, purge=False)
167                     print('Loaded best saved model from {}'.format(
168                         self.model_path))
169 
170     def _wandb_log_predictions(self):
171         "Log prediction samples"
172 
173         pred_log = []
174 
175         for x, y in self.validation_data:
176             try:
177                 pred=self.learn.predict(x)
178             except:
179                 raise FastaiError('Unable to run "predict" method from Learner to log prediction samples.')
180 
181             # scalar -> likely to be a category
182             if not pred[1].shape:
183                 pred_log.append(
184                     wandb.Image(
185                         x.data,
186                         caption='Ground Truth: {}\nPrediction: {}'.format(
187                             y, pred[0])))
188 
189             # most vision datasets have a "show" function we can use
190             elif hasattr(x, "show"):
191                 # log input data
192                 pred_log.append(
193                     wandb.Image(x.data, caption='Input data', grouping=3))
194 
195                 # log label and prediction
196                 for im, capt in ((pred[0], "Prediction"),
197                                  (y, "Ground Truth")):
198                     # Resize plot to image resolution
199                     # from https://stackoverflow.com/a/13714915
200                     my_dpi = 100
201                     fig = plt.figure(frameon=False, dpi=my_dpi)
202                     h, w = x.size
203                     fig.set_size_inches(w / my_dpi, h / my_dpi)
204                     ax = plt.Axes(fig, [0., 0., 1., 1.])
205                     ax.set_axis_off()
206                     fig.add_axes(ax)
207 
208                     # Superpose label or prediction to input image
209                     x.show(ax=ax, y=im)
210                     pred_log.append(wandb.Image(fig, caption=capt))
211                     plt.close(fig)
212 
213             # likely to be an image
214             elif hasattr(y, "shape") and (
215                 (len(y.shape) == 2) or
216                     (len(y.shape) == 3 and y.shape[0] in [1, 3, 4])):
217 
218                 pred_log.extend([
219                     wandb.Image(x.data, caption='Input data', grouping=3),
220                     wandb.Image(pred[0].data, caption='Prediction'),
221                     wandb.Image(y.data, caption='Ground Truth')
222                 ])
223 
224             # we just log input data
225             else:
226                 pred_log.append(wandb.Image(x.data, caption='Input data'))
227 
228             wandb.log({"Prediction Samples": pred_log}, commit=False)
229 
230 
231 class FastaiError(wandb.Error):
232     pass
233
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/wandb/fastai/__init__.py b/wandb/fastai/__init__.py
--- a/wandb/fastai/__init__.py
+++ b/wandb/fastai/__init__.py
@@ -179,7 +179,8 @@
                 raise FastaiError('Unable to run "predict" method from Learner to log prediction samples.')
 
             # scalar -> likely to be a category
-            if not pred[1].shape:
+            # tensor of dim 1 -> likely to be multicategory
+            if not pred[1].shape or pred[1].dim() == 1:
                 pred_log.append(
                     wandb.Image(
                         x.data,

{"golden_diff": "diff --git a/wandb/fastai/__init__.py b/wandb/fastai/__init__.py\n--- a/wandb/fastai/__init__.py\n+++ b/wandb/fastai/__init__.py\n@@ -179,7 +179,8 @@\n                 raise FastaiError('Unable to run \"predict\" method from Learner to log prediction samples.')\n \n             # scalar -> likely to be a category\n-            if not pred[1].shape:\n+            # tensor of dim 1 -> likely to be multicategory\n+            if not pred[1].shape or pred[1].dim() == 1:\n                 pred_log.append(\n                     wandb.Image(\n                         x.data,\n", "issue": "Fastai WandbCallback with multilabel datasets\n* Weights and Biases version: 0.8.9\r\n* Python version: 3.6.9\r\n* Operating System: Linux-4.14.133-88.112.amzn1.x86_64-x86_64-with-Ubuntu-16.04-xenial\r\n\r\n### Description\r\n\r\nI'm trying to see prediction samples for fastai image classification, where the images are multilabel. I'd like  it to be ran on the validation data at the end of each epoch and to see captions of ground truth/predictions.\r\n\r\nWhen I use a multilabel dataset, instead of ground truth/prediction captions for each image, I get a three copies of each image side by side, where each is labeled, \"Input data\", \"Prediction\", \"Ground Truth\" (but not the actual categories)\r\n\r\n### What I Did\r\n\r\n`callback_fns += [partial(WandbCallback, input_type='images', monitor='acc_thresholded')]`\r\n\r\n\n", "before_files": [{"content": "'''\nThis module hooks fast.ai Learners to Weights & Biases through a callback.\nRequested logged data can be configured through the callback constructor.\n\nExamples:\n    WandbCallback can be used when initializing the Learner::\n\n    ```\n        from wandb.fastai import WandbCallback\n        [...]\n        learn = Learner(data, ..., callback_fns=WandbCallback)\n        learn.fit(epochs)\n    ```\n\n    Custom parameters can be given using functools.partial::\n\n    ```\n        from wandb.fastai import WandbCallback\n        from functools import partial\n        [...]\n        learn = Learner(data, ..., callback_fns=partial(WandbCallback, ...))\n        learn.fit(epochs)\n    ```\n\n    Finally, it is possible to use WandbCallback only when starting\n    training. In this case it must be instantiated::\n\n    ```\n        learn.fit(..., callbacks=WandbCallback(learn))\n    ```\n\n    or, with custom parameters::\n\n    ```\n        learn.fit(..., callbacks=WandbCallback(learn, ...))\n    ```\n'''\nimport wandb\nimport fastai\nfrom fastai.callbacks import TrackerCallback\nfrom pathlib import Path\nimport random\ntry:\n    import matplotlib\n    matplotlib.use('Agg')  # non-interactive backend (avoid tkinter issues)\n    import matplotlib.pyplot as plt\nexcept:\n    print('Warning: matplotlib required if logging sample image predictions')\n\n\nclass WandbCallback(TrackerCallback):\n    \"\"\"\n    Automatically saves model topology, losses & metrics.\n    Optionally logs weights, gradients, sample predictions and best trained model.\n\n    Args:\n        learn (fastai.basic_train.Learner): the fast.ai learner to hook.\n        log (str): \"gradients\", \"parameters\", \"all\", or None. Losses & metrics are always logged.\n        save_model (bool): save model at the end of each epoch. It will also load best model at the end of training.\n        monitor (str): metric to monitor for saving best model. None uses default TrackerCallback monitor value.\n        mode (str): \"auto\", \"min\" or \"max\" to compare \"monitor\" values and define best model.\n        input_type (str): \"images\" or None. Used to display sample predictions.\n        validation_data (list): data used for sample predictions if input_type is set.\n        predictions (int): number of predictions to make if input_type is set and validation_data is None.\n        seed (int): initialize random generator for sample predictions if input_type is set and validation_data is None.\n    \"\"\"\n\n    # Record if watch has been called previously (even in another instance)\n    _watch_called = False\n\n    def __init__(self,\n                 learn,\n                 log=\"gradients\",\n                 save_model=True,\n                 monitor=None,\n                 mode='auto',\n                 input_type=None,\n                 validation_data=None,\n                 predictions=36,\n                 seed=12345):\n\n        # Check if wandb.init has been called\n        if wandb.run is None:\n            raise ValueError(\n                'You must call wandb.init() before WandbCallback()')\n\n        # Adapted from fast.ai \"SaveModelCallback\"\n        if monitor is None:\n            # use default TrackerCallback monitor value\n            super().__init__(learn, mode=mode)\n        else:\n            super().__init__(learn, monitor=monitor, mode=mode)\n        self.save_model = save_model\n        self.model_path = Path(wandb.run.dir) / 'bestmodel.pth'\n\n        self.log = log\n        self.input_type = input_type\n        self.best = None\n\n        # Select items for sample predictions to see evolution along training\n        self.validation_data = validation_data\n        if input_type and not self.validation_data:\n            wandbRandom = random.Random(seed)  # For repeatability\n            predictions = min(predictions, len(learn.data.valid_ds))\n            indices = wandbRandom.sample(range(len(learn.data.valid_ds)),\n                                         predictions)\n            self.validation_data = [learn.data.valid_ds[i] for i in indices]\n\n    def on_train_begin(self, **kwargs):\n        \"Call watch method to log model topology, gradients & weights\"\n\n        # Set self.best, method inherited from \"TrackerCallback\" by \"SaveModelCallback\"\n        super().on_train_begin()\n\n        # Ensure we don't call \"watch\" multiple times\n        if not WandbCallback._watch_called:\n            WandbCallback._watch_called = True\n\n            # Logs model topology and optionally gradients and weights\n            wandb.watch(self.learn.model, log=self.log)\n\n    def on_epoch_end(self, epoch, smooth_loss, last_metrics, **kwargs):\n        \"Logs training loss, validation loss and custom metrics & log prediction samples & save model\"\n\n        if self.save_model:\n            # Adapted from fast.ai \"SaveModelCallback\"\n            current = self.get_monitor_value()\n            if current is not None and self.operator(current, self.best):\n                print(\n                    'Better model found at epoch {} with {} value: {}.'.format(\n                        epoch, self.monitor, current))\n                self.best = current\n\n                # Save within wandb folder\n                with self.model_path.open('wb') as model_file:\n                    self.learn.save(model_file)\n\n        # Log sample predictions if learn.predict is available\n        if self.validation_data:\n            try:\n                self._wandb_log_predictions()\n            except FastaiError as e:\n                wandb.termwarn(e.message)\n                self.validation_data = None  # prevent from trying again on next loop\n            except Exception as e:\n                wandb.termwarn(\"Unable to log prediction samples.\\n{}\".format(e))\n                self.validation_data=None  # prevent from trying again on next loop\n\n        # Log losses & metrics\n        # Adapted from fast.ai \"CSVLogger\"\n        logs = {\n            name: stat\n            for name, stat in list(\n                zip(self.learn.recorder.names, [epoch, smooth_loss] +\n                    last_metrics))\n        }\n        wandb.log(logs)\n\n    def on_train_end(self, **kwargs):\n        \"Load the best model.\"\n\n        if self.save_model:\n            # Adapted from fast.ai \"SaveModelCallback\"\n            if self.model_path.is_file():\n                with self.model_path.open('rb') as model_file:\n                    self.learn.load(model_file, purge=False)\n                    print('Loaded best saved model from {}'.format(\n                        self.model_path))\n\n    def _wandb_log_predictions(self):\n        \"Log prediction samples\"\n\n        pred_log = []\n\n        for x, y in self.validation_data:\n            try:\n                pred=self.learn.predict(x)\n            except:\n                raise FastaiError('Unable to run \"predict\" method from Learner to log prediction samples.')\n\n            # scalar -> likely to be a category\n            if not pred[1].shape:\n                pred_log.append(\n                    wandb.Image(\n                        x.data,\n                        caption='Ground Truth: {}\\nPrediction: {}'.format(\n                            y, pred[0])))\n\n            # most vision datasets have a \"show\" function we can use\n            elif hasattr(x, \"show\"):\n                # log input data\n                pred_log.append(\n                    wandb.Image(x.data, caption='Input data', grouping=3))\n\n                # log label and prediction\n                for im, capt in ((pred[0], \"Prediction\"),\n                                 (y, \"Ground Truth\")):\n                    # Resize plot to image resolution\n                    # from https://stackoverflow.com/a/13714915\n                    my_dpi = 100\n                    fig = plt.figure(frameon=False, dpi=my_dpi)\n                    h, w = x.size\n                    fig.set_size_inches(w / my_dpi, h / my_dpi)\n                    ax = plt.Axes(fig, [0., 0., 1., 1.])\n                    ax.set_axis_off()\n                    fig.add_axes(ax)\n\n                    # Superpose label or prediction to input image\n                    x.show(ax=ax, y=im)\n                    pred_log.append(wandb.Image(fig, caption=capt))\n                    plt.close(fig)\n\n            # likely to be an image\n            elif hasattr(y, \"shape\") and (\n                (len(y.shape) == 2) or\n                    (len(y.shape) == 3 and y.shape[0] in [1, 3, 4])):\n\n                pred_log.extend([\n                    wandb.Image(x.data, caption='Input data', grouping=3),\n                    wandb.Image(pred[0].data, caption='Prediction'),\n                    wandb.Image(y.data, caption='Ground Truth')\n                ])\n\n            # we just log input data\n            else:\n                pred_log.append(wandb.Image(x.data, caption='Input data'))\n\n            wandb.log({\"Prediction Samples\": pred_log}, commit=False)\n\n\nclass FastaiError(wandb.Error):\n    pass\n", "path": "wandb/fastai/__init__.py"}], "after_files": [{"content": "'''\nThis module hooks fast.ai Learners to Weights & Biases through a callback.\nRequested logged data can be configured through the callback constructor.\n\nExamples:\n    WandbCallback can be used when initializing the Learner::\n\n    ```\n        from wandb.fastai import WandbCallback\n        [...]\n        learn = Learner(data, ..., callback_fns=WandbCallback)\n        learn.fit(epochs)\n    ```\n\n    Custom parameters can be given using functools.partial::\n\n    ```\n        from wandb.fastai import WandbCallback\n        from functools import partial\n        [...]\n        learn = Learner(data, ..., callback_fns=partial(WandbCallback, ...))\n        learn.fit(epochs)\n    ```\n\n    Finally, it is possible to use WandbCallback only when starting\n    training. In this case it must be instantiated::\n\n    ```\n        learn.fit(..., callbacks=WandbCallback(learn))\n    ```\n\n    or, with custom parameters::\n\n    ```\n        learn.fit(..., callbacks=WandbCallback(learn, ...))\n    ```\n'''\nimport wandb\nimport fastai\nfrom fastai.callbacks import TrackerCallback\nfrom pathlib import Path\nimport random\ntry:\n    import matplotlib\n    matplotlib.use('Agg')  # non-interactive backend (avoid tkinter issues)\n    import matplotlib.pyplot as plt\nexcept:\n    print('Warning: matplotlib required if logging sample image predictions')\n\n\nclass WandbCallback(TrackerCallback):\n    \"\"\"\n    Automatically saves model topology, losses & metrics.\n    Optionally logs weights, gradients, sample predictions and best trained model.\n\n    Args:\n        learn (fastai.basic_train.Learner): the fast.ai learner to hook.\n        log (str): \"gradients\", \"parameters\", \"all\", or None. Losses & metrics are always logged.\n        save_model (bool): save model at the end of each epoch. It will also load best model at the end of training.\n        monitor (str): metric to monitor for saving best model. None uses default TrackerCallback monitor value.\n        mode (str): \"auto\", \"min\" or \"max\" to compare \"monitor\" values and define best model.\n        input_type (str): \"images\" or None. Used to display sample predictions.\n        validation_data (list): data used for sample predictions if input_type is set.\n        predictions (int): number of predictions to make if input_type is set and validation_data is None.\n        seed (int): initialize random generator for sample predictions if input_type is set and validation_data is None.\n    \"\"\"\n\n    # Record if watch has been called previously (even in another instance)\n    _watch_called = False\n\n    def __init__(self,\n                 learn,\n                 log=\"gradients\",\n                 save_model=True,\n                 monitor=None,\n                 mode='auto',\n                 input_type=None,\n                 validation_data=None,\n                 predictions=36,\n                 seed=12345):\n\n        # Check if wandb.init has been called\n        if wandb.run is None:\n            raise ValueError(\n                'You must call wandb.init() before WandbCallback()')\n\n        # Adapted from fast.ai \"SaveModelCallback\"\n        if monitor is None:\n            # use default TrackerCallback monitor value\n            super().__init__(learn, mode=mode)\n        else:\n            super().__init__(learn, monitor=monitor, mode=mode)\n        self.save_model = save_model\n        self.model_path = Path(wandb.run.dir) / 'bestmodel.pth'\n\n        self.log = log\n        self.input_type = input_type\n        self.best = None\n\n        # Select items for sample predictions to see evolution along training\n        self.validation_data = validation_data\n        if input_type and not self.validation_data:\n            wandbRandom = random.Random(seed)  # For repeatability\n            predictions = min(predictions, len(learn.data.valid_ds))\n            indices = wandbRandom.sample(range(len(learn.data.valid_ds)),\n                                         predictions)\n            self.validation_data = [learn.data.valid_ds[i] for i in indices]\n\n    def on_train_begin(self, **kwargs):\n        \"Call watch method to log model topology, gradients & weights\"\n\n        # Set self.best, method inherited from \"TrackerCallback\" by \"SaveModelCallback\"\n        super().on_train_begin()\n\n        # Ensure we don't call \"watch\" multiple times\n        if not WandbCallback._watch_called:\n            WandbCallback._watch_called = True\n\n            # Logs model topology and optionally gradients and weights\n            wandb.watch(self.learn.model, log=self.log)\n\n    def on_epoch_end(self, epoch, smooth_loss, last_metrics, **kwargs):\n        \"Logs training loss, validation loss and custom metrics & log prediction samples & save model\"\n\n        if self.save_model:\n            # Adapted from fast.ai \"SaveModelCallback\"\n            current = self.get_monitor_value()\n            if current is not None and self.operator(current, self.best):\n                print(\n                    'Better model found at epoch {} with {} value: {}.'.format(\n                        epoch, self.monitor, current))\n                self.best = current\n\n                # Save within wandb folder\n                with self.model_path.open('wb') as model_file:\n                    self.learn.save(model_file)\n\n        # Log sample predictions if learn.predict is available\n        if self.validation_data:\n            try:\n                self._wandb_log_predictions()\n            except FastaiError as e:\n                wandb.termwarn(e.message)\n                self.validation_data = None  # prevent from trying again on next loop\n            except Exception as e:\n                wandb.termwarn(\"Unable to log prediction samples.\\n{}\".format(e))\n                self.validation_data=None  # prevent from trying again on next loop\n\n        # Log losses & metrics\n        # Adapted from fast.ai \"CSVLogger\"\n        logs = {\n            name: stat\n            for name, stat in list(\n                zip(self.learn.recorder.names, [epoch, smooth_loss] +\n                    last_metrics))\n        }\n        wandb.log(logs)\n\n    def on_train_end(self, **kwargs):\n        \"Load the best model.\"\n\n        if self.save_model:\n            # Adapted from fast.ai \"SaveModelCallback\"\n            if self.model_path.is_file():\n                with self.model_path.open('rb') as model_file:\n                    self.learn.load(model_file, purge=False)\n                    print('Loaded best saved model from {}'.format(\n                        self.model_path))\n\n    def _wandb_log_predictions(self):\n        \"Log prediction samples\"\n\n        pred_log = []\n\n        for x, y in self.validation_data:\n            try:\n                pred=self.learn.predict(x)\n            except:\n                raise FastaiError('Unable to run \"predict\" method from Learner to log prediction samples.')\n\n            # scalar -> likely to be a category\n            # tensor of dim 1 -> likely to be multicategory\n            if not pred[1].shape or pred[1].dim() == 1:\n                pred_log.append(\n                    wandb.Image(\n                        x.data,\n                        caption='Ground Truth: {}\\nPrediction: {}'.format(\n                            y, pred[0])))\n\n            # most vision datasets have a \"show\" function we can use\n            elif hasattr(x, \"show\"):\n                # log input data\n                pred_log.append(\n                    wandb.Image(x.data, caption='Input data', grouping=3))\n\n                # log label and prediction\n                for im, capt in ((pred[0], \"Prediction\"),\n                                 (y, \"Ground Truth\")):\n                    # Resize plot to image resolution\n                    # from https://stackoverflow.com/a/13714915\n                    my_dpi = 100\n                    fig = plt.figure(frameon=False, dpi=my_dpi)\n                    h, w = x.size\n                    fig.set_size_inches(w / my_dpi, h / my_dpi)\n                    ax = plt.Axes(fig, [0., 0., 1., 1.])\n                    ax.set_axis_off()\n                    fig.add_axes(ax)\n\n                    # Superpose label or prediction to input image\n                    x.show(ax=ax, y=im)\n                    pred_log.append(wandb.Image(fig, caption=capt))\n                    plt.close(fig)\n\n            # likely to be an image\n            elif hasattr(y, \"shape\") and (\n                (len(y.shape) == 2) or\n                    (len(y.shape) == 3 and y.shape[0] in [1, 3, 4])):\n\n                pred_log.extend([\n                    wandb.Image(x.data, caption='Input data', grouping=3),\n                    wandb.Image(pred[0].data, caption='Prediction'),\n                    wandb.Image(y.data, caption='Ground Truth')\n                ])\n\n            # we just log input data\n            else:\n                pred_log.append(wandb.Image(x.data, caption='Input data'))\n\n            wandb.log({\"Prediction Samples\": pred_log}, commit=False)\n\n\nclass FastaiError(wandb.Error):\n    pass\n", "path": "wandb/fastai/__init__.py"}]}