I want to save the model with tta_wrapper and reuse it. However, when reloading the model, a repeat layer problem occurs. What is the solution?

I have loaded the reapeat layer from the tta_wrapper's layer, but the only response is that there is no n in init.

TypeError: init() missing 1 required positional argument: 'n'

Thanks for your tta jobs.
I have a question.

Is this available for models designed in Pytorch?

otherwise we average classes and not probabilities?

AttributeError Traceback (most recent call last)
in ()
2 model = tta_segmentation(model,
3 h_flip=True,
----> 4 h_shift=(-5, 5),
5 #merge='mean'
6 )

C:\Anaconda3\lib\site-packages\tta_wrapper\wrappers.py in tta_segmentation(model, h_flip, v_flip, h_shift, v_shift, rotation, contrast, add, mul, merge)
52 )
53
---> 54 input_shape = (1, *model.input.shape.as_list()[1:])
55
56 inp = Input(batch_shape=input_shape)

C:\Anaconda3\lib\site-packages\keras\engine\base_layer.py in input(self)
782 if len(self._inbound_nodes) > 1:
783 raise AttributeError('Layer ' + self.name +
--> 784 ' has multiple inbound nodes, '
785 'hence the notion of "layer input" '
786 'is ill-defined. '

AttributeError: Layer sequential_1 has multiple inbound nodes, hence the notion of "layer input" is ill-defined. Use get_input_at(node_index) instead.

Below is the traceback and the libraries installed in my conda env.
My suspect is Python 3.5 let me know how to resolve this.
I Believe I need to update the Python 3.5 to 3.7 without disturbing the environment( idk how to do it ) suggestions are welcome

> Traceback (most recent call last):

  File "C:\Users\admin\Miniconda3\envs\machine_learning\lib\site-packages\IPython\core\interactiveshell.py", line 2910, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)

  File "<ipython-input-25-168b6627f61f>", line 1, in <module>
    import tta_wrapper

  File "C:\Users\admin\Miniconda3\envs\machine_learning\lib\site-packages\tta_wrapper\__init__.py", line 1, in <module>
    from .wrappers import tta_segmentation

  File "C:\Users\admin\Miniconda3\envs\machine_learning\lib\site-packages\tta_wrapper\wrappers.py", line 4, in <module>
    from .layers import Repeat, TTA, Merge

  File "C:\Users\admin\Miniconda3\envs\machine_learning\lib\site-packages\tta_wrapper\layers.py", line 58
    raise ValueError(f'Wrong merge type {type}')
                                              ^
SyntaxError: invalid syntax

List of libraries installed in my conda env :

Please update the example in the description, arguments are not relevant

Could you explain why the batch_size must be 1 on inference mode?

This is a nice package to have. So are there any plans to upgrade it for TF Keras? I would like to contribute if you are planning to do so. Let me know your thoughts.