V2: improve CLI for multitask

[hwpang]

Description

• Fix some small bugs related to multitask training and prediction.
• Add a test dataset for single molecule multitask regression
• Add a CLI test for training and predicting multitask
• Improve saving predictions for
  • multi-task
  • multi-task MVE model -> part of v2.1, won't support for now for v2.0.
  • multi-task multiclass -> looking for an example dataset to test this out. -> addressed in V2: Add dataset and tests for multiclass classification #672.

While implementing the cli test for multi-task MVE model, I have encountered bugs related to MVE outside the scope of this PR. I will make another PR to fix it. Edit: single-task MVE is fixed in #664.

Question

• Do we have an example multi-class dataset?

Checklist

• linted with flake8?
• (if appropriate) unit tests added?

[KnathanM]

I now see that the test are failing.

In tests/cli/test_cli_regression_mol.py you have test_train_quick defined twice. The second one uses regression-mve which we are moving to v2.1 so I think you can just delete the second test_train_quick to fix it.

The second error seems a bit more difficult. In tests/cli/test_cli_regression_mol_multitask.py you try to load a multitask model but the model seems to have loc and scale initialized as 1x1 torch tensors, so overwriting them with the 1x12 tensors in the model file throws an error. Maybe the model's ntasks attribute needs to be set sooner so the model will be initialized with a 1x12 loc and scale?

[hwpang]

I took a look into this. The error is the same that I saw previously: when loc is initialized to be a 1 by 1 tensor but the trained model is 1 by n_tasks, it has this error.

I found that the isinstance(loc, float) type check is not working as expected. For some reason, the default loc is recognized as an integer. Changing loc to 0.0 doesn't work. I can change the type check to check for int instead, but it can be confusing.