Недостаточно памяти при выборке изображений с использованием модели авторегрессии TF2.0 в Google Colab

Question

Функция выборки выглядит следующим образом:

def sample_image(batch_size, model):
    image = np.random.choice(4, size=(batch_size, 28, 28, 3))
    for i in range(28):
        for j in range(28):
            for k in range(3):
                _ ,prob_output =  model(tf.Variable(image, dtype=tf.float32, trainable=False))
                prob_output = prob_output.numpy().reshape((batch_size,28,28,3,-1))
                # prob_output = tf.nn.softmax(prob_output, axis=-1)
                # print(prob_output.shape)
                for b in range(batch_size):
                    if k == 0 and b ==0: 
                        print(f'i:{i}, j:{j}, k:{k}')
                    # print(prob_output[b,i,j,k])
                    # prob = tf.nn.softmax(prob_output[b,i,j,k], axis=-1).numpy()
                    # print(prob)
                    # prob /= prob.sum()
                    image[b, i, j, k] = np.random.choice(4, p=prob_output[b,i,j,k])
                del prob_output
                del _

    return image

И когда я запустил этот код для демонстрации изображений в Google Colab, ОЗУ, наконец, стало ООМ, и ядро ​​после этого отключилось. Как избежать OOM при работе на Colab?