def METHOD_NAME():
    with tf.device('gpu:0'):
        model = mnist.keras_model()
        train_dataset = mnist.get_dataset('cpu', 0)
        model.fit(train_dataset, epochs=mnist.EPOCHS, steps_per_epoch=mnist.ITERATIONS)