test_y = np.reshape(test_y, (-1, 3))

for i in range(EPOCHS):
    for i in range(1,hm_data+1):
        train_data = np.load('training_data.npy')

        train = train_data[:-100]
        test = train_data[-100:]

        X = np.array([i[0] for i in train]).reshape(-1,WIDTH,HEIGHT,1)
        Y = [i[1] for i in train]
        Y = np.reshape(Y, (-1, 3))
        print("Shape: ", X.shape)
        print(np.shape(Y))

        test_x = np.array([i[0] for i in test]).reshape(-1,WIDTH,HEIGHT,1)
        test_y = [i[1] for i in test]

        model.fit({'input': X}, {'targets': Y}, n_epoch=1, validation_set=({'input': test_x}, {'targets': test_y}),
            snapshot_step=500, show_metric=True, run_id=MODEL_NAME)

        model.save(MODEL_NAME)