RuntimeError:The size of tensor a (96) must match the size of tensor b (0) at non-singleton dimension 2

[cakmakselin]

I am trying to train the model with 02-train.ipynb on the available CTC data since the pretrained model is not available online. I get the following error:

RuntimeError Traceback (most recent call last)
/Users/.../EmbedSeg/examples/2d/dsb-2018/02-train.ipynb Cell 36 line 1
----> 1 begin_training(train_dataset_dict, val_dataset_dict, model_dict, loss_dict, configs, color_map=new_cmap)

File .../EmbedSeg/criterions/my_loss.py:41, in SpatialEmbLoss.forward(self, prediction, instances, labels, center_images, w_inst, w_var, w_seed, iou, iou_meter)
37 loss = 0
39 for b in range(0, batch_size):
---> 41 spatial_emb = torch.tanh(prediction[b, 0:2]) + xym_s # 2 x h x w #TODO
42 sigma = prediction[b, 2:2 + self.n_sigma] # n_sigma x h x w
43 seed_map = torch.sigmoid(prediction[b, 2 + self.n_sigma:2 + self.n_sigma + 1]) # 1 x h x w

RuntimeError: The size of tensor a (96) must match the size of tensor b (0) at non-singleton dimension 2

Any ideas on how to solve this? Is it alternatively possible to get access to the pretrained models? Thanks!