Error when training model on M3 Max MPS

Please refer to the following thread: Macbook M1 M2 mps acceleration with scVI - #4 by martinkim0