OperatorNotAllowedInGraphError during Transfer Learning

[ccioflan]

Hi!

After creating a conda environment using the provided environment.yml file, followed by additionally installing TensorFlow 2.9.0 as mentioned in the Dockerfile, I tried to run the Jupyter Notebook's cells (put together in a main.py file). When calling the transfer_learning.transfer_learn function, I observed the following error:

File "main.py", line 152, in <module>
    main()
  File "main.py", line 104, in main
    _, model, _ = transfer_learning.transfer_learn(
  File "/path/to/cioflanc/miniconda3/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py", line 302, in wrapper
    return func(*args, **kwargs)
  File "/path/to/cioflanc/few_shot_kws/multilingual_kws/multilingual_kws/embedding/transfer_learning.py", line 76, in transfer_learn
    init_train_ds = audio_dataset.init_single_target(
  File "/path/to/cioflanc/few_shot_kws/multilingual_kws/multilingual_kws/embedding/input_data.py", line 467, in init_single_target
    waveform_ds = waveform_ds.map(self.augment, num_parallel_calls=AUTOTUNE)
  File "/path/to/cioflanc/miniconda3/lib/python3.8/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 1697, in map
    return ParallelMapDataset(
  File "/path/to/cioflanc/miniconda3/lib/python3.8/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 4080, in __init__
    self._map_func = StructuredFunctionWrapper(
  File "/path/to/cioflanc/miniconda3/lib/python3.8/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 3371, in __init__
    self._function = wrapper_fn.get_concrete_function()
  File "/path/to/cioflanc/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 2938, in get_concrete_function
    graph_function = self._get_concrete_function_garbage_collected(
  File "/path/to/cioflanc/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 2906, in _get_concrete_function_garbage_collected
    graph_function, args, kwargs = self._maybe_define_function(args, kwargs)
  File "/path/to/cioflanc/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3213, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/path/to/cioflanc/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3065, in _create_graph_function
    func_graph_module.func_graph_from_py_func(
  File "/path/to/cioflanc/miniconda3/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 986, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/path/to/cioflanc/miniconda3/lib/python3.8/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 3364, in wrapper_fn
    ret = _wrapper_helper(*args)
  File "/path/to/cioflanc/miniconda3/lib/python3.8/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 3299, in _wrapper_helper
    ret = autograph.tf_convert(func, ag_ctx)(*nested_args)
  File "/path/to/cioflanc/miniconda3/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py", line 302, in wrapper
    return func(*args, **kwargs)
  File "/path/to/cioflanc/few_shot_kws/multilingual_kws/multilingual_kws/embedding/input_data.py", line 290, in augment
    self.random_timeshift(audio) if self.max_time_shift_samples > 0 else audio
  File "/path/to/cioflanc/few_shot_kws/multilingual_kws/multilingual_kws/embedding/input_data.py", line 261, in random_timeshift
    if time_shift_amount > 0:
  File "/path/to/cioflanc/miniconda3/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 877, in __bool__
    self._disallow_bool_casting()
  File "/path/to/cioflanc/miniconda3/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 483, in _disallow_bool_casting
    self._disallow_when_autograph_disabled(
  File "/path/to/cioflanc/miniconda3/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 467, in _disallow_when_autograph_disabled
    raise errors.OperatorNotAllowedInGraphError(
tensorflow.python.framework.errors_impl.OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool` is not allowed: AutoGraph is disabled in this function. Try decorating it directly with @tf.function.

Have you noticed this behaviour before? Do you have any suggestions?