SynthSeg Brain Generation fails when randomise_res is set to True

[hvgazula]

what were you trying to do? generating a sample brain with synthseg (using default parameters)

what did you expect will happen? generates a sample without raising any error

what actually happened?

Traceback (most recent call last):
  File "/net/vast-storage/scratch/vast/gablab/hgazula/SynthSeg/scripts/tutorials/1-generation_visualisation.py", line 28, in <module>
    im, lab = brain_generator.generate_brain()
  File "/om2/user/hgazula/SynthSeg/SynthSeg/brain_generator.py", line 324, in generate_brain
    (image, labels) = next(self.brain_generator)
  File "/om2/user/hgazula/SynthSeg/SynthSeg/brain_generator.py", line 319, in _build_brain_generator
    [image, labels] = self.labels_to_image_model.predict(model_inputs)
  File "/om/user/hgazula/venvs/nobrainer_satra/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/tmp/__autograph_generated_filescynv096.py", line 15, in tf__predict_function
    retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
  File "/tmp/__autograph_generated_filedv9fe6xu.py", line 38, in tf__call
    ag__.if_stmt(ag__.not_(ag__.ld(self).add_batchsize), if_body, else_body, get_state, set_state, ('mask', 'self.min_res_tens', 'shape'), 3)
  File "/tmp/__autograph_generated_filedv9fe6xu.py", line 28, in else_body
    ag__.ld(self).min_res_tens = ag__.converted_call(ag__.ld(tf).tile, (ag__.converted_call(ag__.ld(tf).expand_dims, (ag__.ld(self).min_res_tens, 0), None, fscope), ag__.ld(tile_shape)), None, fscope)
ValueError: in user code:

    File "/om/user/hgazula/venvs/nobrainer_satra/lib/python3.10/site-packages/keras/src/engine/training.py", line 2416, in predict_function  *
        return step_function(self, iterator)
    File "/om/user/hgazula/venvs/nobrainer_satra/lib/python3.10/site-packages/keras/src/engine/training.py", line 2401, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/om/user/hgazula/venvs/nobrainer_satra/lib/python3.10/site-packages/keras/src/engine/training.py", line 2389, in run_step  **
        outputs = model.predict_step(data)
    File "/om/user/hgazula/venvs/nobrainer_satra/lib/python3.10/site-packages/keras/src/engine/training.py", line 2357, in predict_step
        return self(x, training=False)
    File "/om/user/hgazula/venvs/nobrainer_satra/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler
        raise e.with_traceback(filtered_tb) from None
    File "/tmp/__autograph_generated_filedv9fe6xu.py", line 38, in tf__call
        ag__.if_stmt(ag__.not_(ag__.ld(self).add_batchsize), if_body, else_body, get_state, set_state, ('mask', 'self.min_res_tens', 'shape'), 3)
    File "/tmp/__autograph_generated_filedv9fe6xu.py", line 28, in else_body
        ag__.ld(self).min_res_tens = ag__.converted_call(ag__.ld(tf).tile, (ag__.converted_call(ag__.ld(tf).expand_dims, (ag__.ld(self).min_res_tens, 0), None, fscope), ag__.ld(tile_shape)), None, fscope)

    ValueError: Exception encountered when calling layer 'sample_resolution' (type SampleResolution).
    
    in user code:
    
        File "/om2/user/hgazula/SynthSeg/ext/lab2im/layers.py", line 608, in call  *
            self.min_res_tens = tf.tile(tf.expand_dims(self.min_res_tens, 0), tile_shape)
    
        ValueError: Shape must be rank 3 but is rank 2 for '{{node model/sample_resolution/Tile}} = Tile[T=DT_FLOAT, Tmultiples=DT_INT32](model/sample_resolution/ExpandDims, model/sample_resolution/concat)' with input shapes: [1,?,3], [2].
    
    
    Call arguments received by layer 'sample_resolution' (type SampleResolution):
      • inputs=tf.Tensor(shape=(None, 113, 1), dtype=float32)
      • kwargs={'training': 'False'}

Can you replicate the behavior? If yes, how? see github.com/nobrainer_training_scripts/1.2.0/scripts/train/synthseg.py