Проблемы с использованием предварительно обученной модели от niftynet

Я хочу сделать сегментацию на 2 класса, используя модель dense_vnet, доступную на niftynet, которая изначально выполняет сегментацию на 9 классов.

Я попытался переобучить только последний слой, внеся изменения в файл конфигурации в соответствии с этим предложением: КАК настроить предварительно обученную модель niftynet для пользовательских данных

vars_to_restore = ^((?!DenseVNet/(skip_conv|fin_conv)).)*$

num_classes = 2

ошибка:

tensorflow.python.framework.errors_impl.InvalidArgumentError: Assign requires shapes of both tensors to match. lhs shape= [2] rhs shape= [9] [[{{node save/Assign_8}} = Assign[T=DT_FLOAT, class=["loc:@DenseVNet/conv/conv/b"], use_locking=true, validate_shape=true, device = "/job:localhost/replica:0/task:0/device:CPU:0"](DenseVNet/conv/conv/b, save/RestoreV2:8)]]