Wavenet not running fully on Mac

[zannads]

When training the Wavenet model on the Mac with Metal Acceleration (device 'mps:0'), the training works but at the end when computing the final tensor it fails.
It works on CPU and on CUDA however.

This is an important Issue because it allows a 5x 6x acceleration.

[zannads]

Wavenet is not running on the Mac GPU.
It runs the training and it finishes, however, when adjusting the bias at the end of the iteration the MultiOutputWrapper function throws an error.
In particular, in file multioutput.py, function _get_args_kwargs_by_output, at row 110, inside apply_to_collection(args, Tensor,.. the error is

return tuple(function(x, *args, **kwargs) for x in data)
IndexError: Dimension specified as -1 but tensor has no dimensions

Probably something that has to do on some internal definition of the metrics of the multioutput wrapper (?).

No problem appear when running on the cpu device!