Potential bug in comet-score

[mahmoudaymo]

I have trained a model using MSE and the used the trained model to predict the scores and stds.

the output of comet-score is a bit weird. The same score and std every to two segments as follow:

comet-score -s src.txt -t tgt.txt --model mdel.ckpt --mc_dropout 100 --refless true --batch_size 1 --gpus 1

Predicting: 100%|██████████| 100/100 [07:10<00:00,  4.30s/it]
Segment 0       score: 0.8868   variance: 0.0315
Segment 1       score: 0.8868   variance: 0.0315
Segment 2       score: 0.9553   variance: 0.0367
Segment 3       score: 0.9553   variance: 0.0367
Segment 4       score: 0.9335   variance: 0.0292
Segment 5       score: 0.9335   variance: 0.0292
Segment 6       score: 0.8740   variance: 0.0341
Segment 7       score: 0.8740   variance: 0.0341
Segment 8       score: 0.8326   variance: 0.0426
Segment 9       score: 0.8326   variance: 0.0426
Segment 10      score: 0.9717   variance: 0.0578
Segment 11      score: 0.9717   variance: 0.0578

In addition, when running the command comet-score with batch_size > 1, I get the following error:

Predicting: 0it [00:00, ?it/s]Traceback (most recent call last):
  File "/usr/local/bin/comet-score", line 8, in <module>
    sys.exit(score_command())
  File "/uncertainties_MT_eval/comet/cli/score.py", line 133, in score_command
    mean_scores, std_scores, sys_score = model.predict(
  File "/uncertainties_MT_eval/comet/models/base.py", line 518, in predict
    predictions = trainer.predict(
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/trainer/trainer.py", line 631, in predict
    results = self._run(model)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/trainer/trainer.py", line 756, in _run
    self.dispatch()
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/trainer/trainer.py", line 795, in dispatch
    self.accelerator.start_predicting(self)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/accelerators/accelerator.py", line 102, in start_predicting
    self.training_type_plugin.start_predicting(trainer)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 152, in start_predicting
    self._results = trainer.run_stage()
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/trainer/trainer.py", line 806, in run_stage
    return self.run_predict()
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/trainer/trainer.py", line 1083, in run_predict
    self.predict_loop.predict_step(batch, batch_idx, dataloader_idx)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/trainer/predict_loop.py", line 111, in predict_step
    predictions = self.trainer.accelerator.predict_step(args)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/accelerators/accelerator.py", line 265, in predict_step
    return self.training_type_plugin.predict_step(*args)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 167, in predict_step
    return self.lightning_module.predict_step(*args, **kwargs)
  File "/uncertainties_MT_eval/comet/models/base.py", line 372, in predict_step
    mcd_outputs[i,:] = outputs["score"].view(-1)
RuntimeError: The expanded size of the tensor (2) must match the existing size (10) at non-singleton dimension 0.  Target sizes: [2].  Tensor sizes: [10]

Could you provide any guidance or help on the two issues.