Отрисовка ax.imshow() в цикле for

Имеется вполне рабочая функция реализации Игры Жизни Конвея:

torch_update(Z)

Написана на PyTorch

И код отрисовки, который должен рисовать обновляемый тензор выхода вышеупомянутой функции. Вот он:

%matplotlib notebook
plt.ion()

#initialize game field
Z = np.random.choice([0,1],size=(100,100))
Z = torch.from_numpy(Z).type(torch.FloatTensor)

fig = plt.figure()
ax = fig.add_subplot(111)
fig.show()

for _ in range(100):
    #update
    Z = torch_update(Z)

    ax.clear()
    ax.imshow(Z.numpy(),cmap='gray')
    fig.canvas.draw()

При его запуске получается пустой выход, хотя если засунуть в цикл принт каждой итерации, будет видно, что выход каждый раз меняется.

Подскажите, в чем ошибка в коде отрисовки и как это исправить.

P.S. при работе в Jupiter Notebook локально все ок. Проблема возникает в Google Colab