Add training notebook link.

README.md

@@ -64,6 +64,7 @@ The following code gets the data and preprocesses/augments the data.
 
 ```python
 from datasets import load_dataset
+from super_image.data import EvalDataset, TrainDataset, augment_five_crop
 
 augmented_dataset = load_dataset('eugenesiow/Div2k', 'bicubic_x4', split='train')\
     .map(augment_five_crop, batched=True, desc="Augmenting Dataset")                                # download and augment the data with the five_crop method
@@ -89,11 +90,13 @@ trainer = Trainer(
     model=model,                         # the instantiated model to be trained
     args=training_args,                  # training arguments, defined above
     train_dataset=train_dataset,         # training dataset
-    eval_dataset=val_dataset             # evaluation dataset
+    eval_dataset=eval_dataset            # evaluation dataset
 )
 
 trainer.train()
 ```
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/eugenesiow/super-image-notebooks/blob/master/notebooks/Train_super_image_Models.ipynb "Open in Colab")
 ## Evaluation results
 The evaluation metrics include [PSNR](https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio#Quality_estimation_with_PSNR) and [SSIM](https://en.wikipedia.org/wiki/Structural_similarity#Algorithm).