TF .pb wont save in Colab local

[BigJack325]

To save in Colab, we need to add some lines if we train with TPU because it usually requires to be saved to a Bucket. To overcome the use of a Bucket, add the SaveOptions as follows to the fit() function:

        if export_model:
            save_locally = tf.saved_model.SaveOptions(experimental_io_device='/job:localhost')
            self.model.save(export_model, options=save_locally)
            print('Model exported')