  import sys
  import time
  
  from collections import OrderedDict
  from typing import Dict, List, Optional, Tuple, TYPE_CHECKING, Union
  
  import numpy as np
  
  from lib.serializer import get_serializer
  from lib.model.nn_blocks import set_config as set_nnblock_config
 
</s> remove     def output_shapes(self) -> List[List[Tuple[int, int, int]]]:
        """ list: A list of list of shape tuples for the outputs of the model with the batch
        dimension removed. The outer list contains 2 sub-lists (one for each side "a" and "b").
        The inner sub-lists contain the output shapes for that side. """
        shapes: List[Tuple[int, int, int]] = [tuple(K.int_shape(output)[-3:])  # type: ignore
                                              for output in self.model.outputs]
        return [shapes[:len(shapes) // 2], shapes[len(shapes) // 2:]]
 </s> add     def input_shapes(self) -> List[Tuple[None, int, int, int]]:
        """ list: A flattened list corresponding to all of the inputs to the model. """
        shapes = [cast(Tuple[None, int, int, int], K.int_shape(inputs))
                  for inputs in self.model.inputs]
        return shapes

    @property
    def output_shapes(self) -> List[Tuple[None, int, int, int]]:
        """ list: A flattened list corresponding to all of the outputs of the model. """
        shapes = [cast(Tuple[None, int, int, int], K.int_shape(output))
                  for output in self.model.outputs]
        return shapes </s> remove         self._model_input_size = max(img[1] for img in model.model.input_shape)
 </s> add         self._model_input_size = max(img[1] for img in model.input_shapes) </s> remove         self._process_size = max(img[1]
                                 for img in model.model.input_shape + model.model.output_shape)
        self._output_sizes = [shape[0] for shape in model.output_shapes[0] if shape[-1] != 1]
 </s> add         self._process_size = max(img[1] for img in model.input_shapes + model.output_shapes)
        self._output_sizes = [shape[1] for shape in model.output_shapes if shape[-1] != 1]

from typing import cast, Dict, List, Optional, Tuple, TYPE_CHECKING, Union
  import sys
  import time
  
  from collections import OrderedDict
from typing import cast, Dict, List, Optional, Tuple, TYPE_CHECKING, Union
  
  import numpy as np
  
  from lib.serializer import get_serializer
  from lib.model.nn_blocks import set_config as set_nnblock_config
keep keep keep keep replace keep keep keep keep keep
 bugfix: Get correct output size for learn mask 
plugins/train/model/_base/model.py