update generator to inherit from structures_to_learn size
" \n",
"from_numpy_generator_params = (\n",
" (tf.string, tf.int32, tf.float64), \n",
- " (tf.TensorShape(()), tf.TensorShape([512, 512, 1]), tf.TensorShape([512, 512, 41]))\n",
+ " (\n",
+ " tf.TensorShape(()), \n",
+ " tf.TensorShape([512, 512, 1]), \n",
+ " tf.TensorShape([512, 512, len(structures_to_learn)]))\n",
")\n",
"\n",
"from_numpy_dataset = tf.data.Dataset.from_generator(\n",
