[REVIEW] Support dict output for models

[VibhuJawa]

This PR adds support for models that give dictionary output.

Other changes include:

• Bump up CF version
• Add worker address (to fix bug on SLURM)

Main takeaway is support below model:

class DummyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, batch):
        output_size = len(batch)
        return {
            "a": torch.ones(output_size, device="cuda") * 1,
            "b": torch.ones(output_size, device="cuda") * 2,
        }

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[VibhuJawa]

/okay to test

[VibhuJawa]

/okay to test

[VibhuJawa]

/okay to test

[praateekmahajan]

Conceptually makes sense, happy to approve since it's critical functionality tho would recommend

1. adding more checks / documentation (about behavior / expectation of dict)
2. type hint / docstring in get_model_output(..)
3. deciding signature between model_output_col / model_output_cols (deprecating one of them)
4. documenting behaviour of meta and expectation of dtypes when dict (same as point 1)

[VibhuJawa]

@praateekmahajan ,

adding more checks / documentation (about behavior / expectation of dict)
Added more checks, please take a look.

type hint / docstring in get_model_output(..)
Done

deciding signature between model_output_col / model_output_cols (deprecating one of them)
Created a warning, will deprecate next release.

documenting behaviour of meta and expectation of dtypes when dict (same as point 1)

Done.

Please take another look. Thanks

[praateekmahajan]

nit, but the output_type isn't Any right? It's Union[ModelOutputType, Dict[str, ModelOutputType]

[praateekmahajan]

LGTM!

[VibhuJawa]

/okay to test