Fix output error occurring due to check if it is a dict or not (#1742)

* fix pred output type

* remove commented out lines

* Apply suggestions from code review

Co-authored-by: Karl Higley <[email protected]>

Co-authored-by: Sara Rabhi <[email protected]>
Co-authored-by: Karl Higley <[email protected]>

Author: 3 people

nvtabular/inference/triton/model/model_pt.py

@@ -157,12 +157,7 @@ def execute(self, requests):
 
                 # Call forward function to get the predictions
                 # Forward function should return a dict with the "predictions" bucket
-                out = self.model(input_dict, training=False)
-                if not isinstance(out, dict):
-                    raise ValueError("output of the forward function should be a dict")
-
-                # Get the predictions from the out
-                pred = out.get("predictions")
+                pred = self.model(input_dict, training=False)
                 if pred is None:
                     raise KeyError(
                         "output of the forward function should have a bucket named as predictions"