Python kernel dies on Jupyter Notebook with tensorflow 2

I installed tensorflow 2 on my mac using conda according these instructions:

conda create -n tf2 tensorflow

Then I installed ipykernel to add this new environment to my jupyter notebook kernels as follows:

conda activate tf2
conda install ipykernel
python -m ipykernel install --user --name=tf2

That seemed to work well, I am able to see my tf2 environment on my jupyter notebook kernels.

Then I tried to run the simple MNIST example to check if all was working properly and I when I execute this line of code:

model.fit(x_train, y_train, epochs=5)

The kernel of my jupyter notebook dies without more information.

I executed the same code on my terminal via python mnist_test.py and also via ipython (command by command) and I don’t have any issues, which let’s me assume that my tensorflow 2 is correctly installed on my conda environment.

Any ideas on what went wrong during the install?

Versions:

python==3.7.5
tensorboard==2.0.0
tensorflow==2.0.0
tensorflow-estimator==2.0.0
ipykernel==5.1.3
ipython==7.10.2
jupyter==1.0.0
jupyter-client==5.3.4
jupyter-console==5.2.0
jupyter-core==4.6.1

Here I put the complete script as well as the STDOUT of the execution:

import tensorflow as tf
import matplotlib.pyplot as plt
import seaborn as sns

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()

x_train, x_test = x_train / 255.0, x_test / 255.0

nn_model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation='softmax')
])

nn_model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

nn_model.fit(x_train, y_train, epochs=5)

nn_model.evaluate(x_test,  y_test, verbose=2)

(tf2) ➜ tensorflow2 python mnist_test.py 2020-01-03 10:46:10.854619:
I tensorflow/core/platform/cpu_feature_guard.cc:145] This TensorFlow
binary is optimized with Intel(R) MKL-DNN to use the following CPU
instructions in performance critical operations: SSE4.1 SSE4.2 AVX
AVX2 FMA To enable them in non-MKL-DNN operations, rebuild TensorFlow
with the appropriate compiler flags. 2020-01-03 10:46:10.854860: I
tensorflow/core/common_runtime/process_util.cc:115] Creating new
thread pool with default inter op setting: 8. Tune using
inter_op_parallelism_threads for best performance. Train on 60000
samples Epoch 1/5 60000/60000 [==============================] – 6s
102us/sample – loss: 0.3018 – accuracy: 0.9140 Epoch 2/5 60000/60000
[==============================] – 6s 103us/sample – loss: 0.1437 –
accuracy: 0.9571 Epoch 3/5 60000/60000
[==============================] – 6s 103us/sample – loss: 0.1054 –
accuracy: 0.9679 Epoch 4/5 60000/60000
[==============================] – 6s 103us/sample – loss: 0.0868 –
accuracy: 0.9729 Epoch 5/5 60000/60000
[==============================] – 6s 103us/sample – loss: 0.0739 –
accuracy: 0.9772 10000/1 – 1s – loss: 0.0359 – accuracy: 0.9782 (tf2)
➜ tensorflow2