InvalidArgumentError: Inputs to operation loss/Classification_loss/logistic_loss/Select of type Sele

Epoch 1/100
13097/13097 [==============================] - 5262s 402ms/step - loss: 2.8854 - acc: 0.8722 - iou: 0.2555 - val_loss: 5.3026 - val_acc: 0.9111 - val_iou: 0.2434

Epoch 00001: saving model to seg_hrnet-01-5.3026-0.9111-0.2434.hdf5
Epoch 2/100
---------------------------------------------------------------------------
InvalidArgumentError: Inputs to operation loss/Classification_loss/logistic_loss/Select of type Select must have the same size and shape.  Input 0: [2,1,512,512] != input 1: [1,1,512,512]

问题：出现batch维度的不匹配，造成原因，图片/batchsize不能整除。

我这里的情况是batch=2，因此最后不能整除剩余1张图片，导致此报错。

也就是要修改一下batch_generater,在其内部，若是检测到索引超过了图片数，则从头开始取。

我的示例代码如下：

# 获取batch data
def batch_generator(train_images_paths, batch_size, flag):
    while 1:
        for i in range(0, len(train_images_paths), batch_size):
            idx_start = 0 if (i + batch_size) > len(train_images_paths) else i
            idx_end = idx_start + batch_size
            if flag == 'train':
                images, slic_seg, gts = read_train_img(train_images_paths[idx_start: idx_end])
            else:
                images, slic_seg, gts = read_test_img(train_images_paths[idx_start: idx_end])
            yield ([images, slic_seg], [gts])