Probabilities do not sum to 1

( Note I am running this on the previous version as I cannot use the latest version, as per #3 ) When I try run stepmania_note_generator on a fully retrained model, I get the below error:

Generating notes for breakitoff
-----------------------------------------

0.018300232973530228
Traceback (most recent call last):
  File "C:\dev\StepCOVNet-ef61b36bf2a30bfacb3fdd49a3a613e38d5ca548\stepmania_note_generator.py", line 154, in <module>
    stepmania_note_generator(args.input,
  File "C:\dev\StepCOVNet-ef61b36bf2a30bfacb3fdd49a3a613e38d5ca548\stepmania_note_generator.py", line 126, in stepmania_note_generator
    generate_notes(output_path, tmp_dir, stepcovnet_model, verbose_int)
  File "C:\dev\StepCOVNet-ef61b36bf2a30bfacb3fdd49a3a613e38d5ca548\stepmania_note_generator.py", line 88, in generate_notes
    pred_arrows = inference_executor.execute(input_data=inference_input)
  File "C:\dev\StepCOVNet-ef61b36bf2a30bfacb3fdd49a3a613e38d5ca548\stepcovnet\executor\InferenceExecutor.py", line 43, in execute
    encoded_arrow = np.random.choice(NUM_ARROW_TYPES, 1, p=binary_arrow_prob)[0]
  File "mtrand.pyx", line 939, in numpy.random.mtrand.RandomState.choice
ValueError: probabilities do not sum to 1

As you can see I added a print(sum(binary_arrow_prob)) in InferenceExecutor.py and it's clear that it's not even close to 1 and therefore I can't just correct with rounding. I did initially try that just in case but it results in like 9000 notes for a 1 minute song lol.

Any ideas? I've spent about 2 days trying to get this working initially, then 2 days running the training and I'm starting to lose hope on getting this working.

Many Thanks

cpuguy96 / StepCOVNet

Probabilities do not sum to 1 #4