[AIR] Predictor call_model API for unsupported output types

[amogkam]

Signed-off-by: Amog Kamsetty [email protected]

Adds a Public call_model API to TorchPredictor and TensorflowPredictor.

We also remove default support for List outputs, as it's ambiguous if the list represents the batch, or different output columns.

Supported output types are tensor and Dict[str, tensor].

To handle models with unsupported output types, users are expected to override call_model.

Example:

import torch

class SSDPredictor(TorchPredictor):
    def call_model(self, tensor):
        output = super().call_model(tensor)
        output = {k: torch.stack([d[k] for d in output]) for k in output[0]}
        return output

Documentation will be added if this API looks good.

Why are these changes needed?

Related issue number

Checks

• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[clarkzinzow]

LGTM, nice examples in the call_model() docstrings!