Skip to content

Commit

Permalink
Check for tensor
Browse files Browse the repository at this point in the history
  • Loading branch information
adelavega committed Nov 15, 2022
1 parent 4a6543c commit d92c637
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions pliers/extractors/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,9 +118,11 @@ def _postprocess(self, out):
# If output is a dict and no output key, return all keys
if isinstance(out, dict):
out = np.vstack(list(out.values())).T

elif isinstance(out, tf.Tensor):
out = out.numpy()

# Always squeeze last dimension if it is 1
out = out.numpy().squeeze()
out = out.squeeze()

if self.transform_out:
out = self.transform_out(out)
Expand Down

0 comments on commit d92c637

Please sign in to comment.