Repository: voicepaw/so-vits-svc-fork

Issue:
MPS Support
I've read some issues about mps of pytorch, it turns out that currently mps doesn't support complex types (like 1+2j). But I think svc requires complex types. One of the current solution is adding a.to("cpu") before the operations which are not supported and a.to("mps") after that.

Could this be a temporary workaround for an M1 version than can train?
Once pytorch supports all these operations, though, these added codes should be removed.

```
$ svc train
[16:23:40] INFO     [16:23:40] Version: 3.1.10                                                                                                                                      __main__.py:21
[16:23:41] INFO     [16:23:41] Created a temporary directory at /var/folders/26/3xzsh5qs4zd3_0nnhxsnckrr0000gn/T/tmp3tlrptr6                                                    instantiator.py:21
           INFO     [16:23:41] Writing /var/folders/26/3xzsh5qs4zd3_0nnhxsnckrr0000gn/T/tmp3tlrptr6/_remote_module_non_scriptable.py                                            instantiator.py:76
INFO: GPU available: True (mps), used: True
[16:23:43] INFO     [16:23:43] GPU available: True (mps), used: True                                                                                                               rank_zero.py:48
INFO: TPU available: False, using: 0 TPU cores
           INFO     [16:23:43] TPU available: False, using: 0 TPU cores                                                                                                            rank_zero.py:48
INFO: IPU available: False, using: 0 IPUs
           INFO     [16:23:43] IPU available: False, using: 0 IPUs                                                                                                                 rank_zero.py:48
INFO: HPU available: False, using: 0 HPUs
           INFO     [16:23:43] HPU available: False, using: 0 HPUs                                                                                                                 rank_zero.py:48
           WARNING  [16:23:43] /opt/miniconda3/envs/vox/lib/python3.10/site-packages/so_vits_svc_fork/modules/synthesizers.py:81: UserWarning: Unused arguments: {'n_layers_q': 3, warnings.py:109
                    'use_spectral_norm': False}                                                                                                                                                   
                      warnings.warn(f"Unused arguments: {kwargs}")                                                                                                                                
                                                                                                                                                                                                  
           INFO     [16:23:43] Decoder type: hifi-gan                                                                                                                          synthesizers.py:100
[16:23:44] WARNING  [16:23:44] /opt/miniconda3/envs/vox/lib/python3.10/site-packages/so_vits_svc_fork/utils.py:187: UserWarning: Keys not found in checkpoint state                warnings.py:109
                    dict:['emb_g.weight']                                                                                                                                                         
                      warnings.warn(f"Keys not found in checkpoint state dict:" f"{not_in_from}")                                                                                                 
                                                                                                                                                                                                  
           INFO     [16:23:44] Loaded checkpoint 'logs/44k/G_0.pth' (iteration 0)                                                                                                     utils.py:243
           INFO     [16:23:44] Loaded checkpoint 'logs/44k/D_0.pth' (iteration 0)                                                                                                     utils.py:243
INFO: 
  | Name  | Type                     | Params
---------------------------------------------------
0 | net_g | SynthesizerTrn           | 45.2 M
1 | net_d | MultiPeriodDiscriminator | 46.7 M
---------------------------------------------------
91.9 M    Trainable params
0         Non-trainable params
91.9 M    Total params
367.617   Total estimated model params size (MB)
           INFO     [16:23:44]                                                                                                                                                 model_summary.py:83
                      | Name  | Type                     | Params                                                                                                                                 
                    ---------------------------------------------------                                                                                                                           
                    0 | net_g | SynthesizerTrn           | 45.2 M                                                                                                                                 
                    1 | net_d | MultiPeriodDiscriminator | 46.7 M                                                                                                                                 
                    ---------------------------------------------------                                                                                                                           
                    91.9 M    Trainable params                                                                                                                                                    
                    0         Non-trainable params                                                                                                                                                
                    91.9 M    Total params                                                                                                                                                        
                    367.617   Total estimated model params size (MB)                                                                                                                              
Sanity Checking: 0it [00:00, ?it/s]           WARNING  [16:23:44] /opt/miniconda3/envs/vox/lib/python3.10/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:430: PossibleUserWarning: The          warnings.py:109
                    dataloader, val_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 10 which                  
                    is the number of cpus on this machine) in the `DataLoader` init to improve performance.                                                                                       
                      rank_zero_warn(                                                                                                                                                             
                                                                                                                                                                                                  
           WARNING  [16:23:44] /opt/miniconda3/envs/vox/lib/python3.10/site-packages/torch/_utils.py:776: UserWarning: TypedStorage is deprecated. It will be removed in the       warnings.py:109
                    future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage                        
                    directly, use tensor.untyped_storage() instead of tensor.storage()                                                                                                            
                      return self.fget.__get__(instance, owner)()                                                                                                                                 
                                                                                                                                                                                                  
Sanity Checking DataLoader 0:   0%|                                                                                                                                         | 0/2 [00:00<?, ?it/s]           WARNING  [16:23:44] /opt/miniconda3/envs/vox/lib/python3.10/site-packages/so_vits_svc_fork/f0.py:221: UserWarning: MPS: nonzero op is supported natively starting from  warnings.py:109
                    macOS 13.0. Falling back on CPU. This may have performance implications. (Triggered internally at                                                                             
                    /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/mps/operations/Indexing.mm:218.)                                                                              
                      f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * (f0_bin - 2) / (                                                                                                   
                                                                                                                                                                                                  
           WARNING  [16:23:44] /opt/miniconda3/envs/vox/lib/python3.10/site-packages/so_vits_svc_fork/f0.py:228: UserWarning: MPS: no support for int64 min/max ops, casting it to warnings.py:109
                    int32 (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/mps/operations/ReduceOps.mm:1271.)                                             
                      assert f0_coarse.max() <= 255 and f0_coarse.min() >= 1, (                                                                                                                   
                                                                                                                                                                                                  
           WARNING  [16:23:44] /opt/miniconda3/envs/vox/lib/python3.10/site-packages/so_vits_svc_fork/modules/attentions.py:391: UserWarning: MPS: The constant padding of more    warnings.py:109
                    than 3 dimensions is not currently supported natively. It uses View Ops default implementation to run. This may have performance implications. (Triggered                     
                    internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/mps/operations/Pad.mm:393.)                                                                     
                      x = F.pad(x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, 1]]))                                                                                                   
                                                                                                                                                                                                  
           WARNING  [16:23:44] /opt/miniconda3/envs/vox/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:25: UserWarning: The operator 'aten::_weight_norm_interface' is warnings.py:109
                    not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at                             
                    /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)                                                                                              
                      return _weight_norm(v, g, self.dim)                                                                                                                                         
                                                                                                                                                                                                  
           WARNING  [16:23:44] /opt/miniconda3/envs/vox/lib/python3.10/site-packages/torch/nn/functional.py:3929: UserWarning: MPS: passing scale factor to upsample ops is        warnings.py:109
                    supported natively starting from macOS 13.0. Falling back on CPU. This may have performance implications. (Triggered internally at                                            
                    /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/mps/operations/UpSample.mm:233.)                                                                              
                      return torch._C._nn.upsample_nearest1d(input, output_size, scale_factors)                                                                                                   
                                                                                                                                                                                                  
           WARNING  [16:23:44] /opt/miniconda3/envs/vox/lib/python3.10/site-packages/so_vits_svc_fork/modules/decoders/hifigan/_models.py:86: UserWarning: torch.cumsum supported  warnings.py:109
                    by MPS on MacOS 13+, please upgrade (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/mps/operations/UnaryOps.mm:264.)                 
                      tmp_over_one = torch.cumsum(rad_values, 1) % 1                                                                                                                              
                                                                                                                                                                                                  
[16:23:45] WARNING  [16:23:45] /opt/miniconda3/envs/vox/lib/python3.10/site-packages/torch/functional.py:641: UserWarning: stft with return_complex=False is deprecated. In a      warnings.py:109
                    future pytorch release, stft will return complex tensors for all inputs, and return_complex=False will raise an error.                                                        
                    Note: you can still call torch.view_as_real on the complex output to recover the old return format. (Triggered internally at                                                  
                    /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/SpectralOps.cpp:867.)                                                                                         
                      return _VF.stft(input, n_fft, hop_length, win_length, window,  # type: ignore[attr-defined]                                                                                 
                                                                                                                                                                                                  
[16:23:46] WARNING  [16:23:46] /opt/miniconda3/envs/vox/lib/python3.10/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:430: PossibleUserWarning: The          warnings.py:109
                    dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 10 which                
                    is the number of cpus on this machine) in the `DataLoader` init to improve performance.                                                                                       
                      rank_zero_warn(                                                                                                                                                             
                                                                                                                                                                                                  
           WARNING  [16:23:46] /opt/miniconda3/envs/vox/lib/python3.10/site-packages/lightning/pytorch/loops/fit_loop.py:280: PossibleUserWarning: The number of training batches  warnings.py:109
                    (10) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training                     
                    epoch.                                                                                                                                                                        
                      rank_zero_warn(                                                                                                                                                             
                                                                                                                                                                                                  
Training: 0it [00:00, ?it/s]           INFO     [16:23:46] Setting current epoch to 0                                                                                                                             train.py:198
           INFO     [16:23:46] Setting total batch idx to 0                                                                                                                           train.py:213
           INFO     [16:23:46] Setting global step to 0                                                                                                                               train.py:203
Epoch 0:   0%|                                                                                                                                                             | 0/10 [00:00<?, ?it/s]libc++abi: terminating with uncaught exception of type c10::Error: Unsupported type byte size: ComplexFloat
Exception raised from getGatherScatterScalarType at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/mps/operations/View.mm:758 (most recent call first):
frame #0: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > const&) + 92 (0x103f8d2b8 in libc10.dylib)
frame #1: at::native::mps::getGatherScatterScalarType(at::Tensor const&) + 304 (0x16692266c in libtorch_cpu.dylib)
frame #2: invocation function for block in at::native::mps::gatherViewTensor(at::Tensor const&, at::Tensor&) + 128 (0x1669241bc in libtorch_cpu.dylib)
frame #3: _dispatch_client_callout + 20 (0x1bee501b4 in libdispatch.dylib)
frame #4: _dispatch_lane_barrier_sync_invoke_and_complete + 56 (0x1bee5f414 in libdispatch.dylib)
frame #5: at::native::mps::gatherViewTensor(at::Tensor const&, at::Tensor&) + 888 (0x166922d54 in libtorch_cpu.dylib)
frame #6: at::native::mps::mps_copy_(at::Tensor&, at::Tensor const&, bool) + 3096 (0x16687a47c in libtorch_cpu.dylib)
frame #7: at::native::copy_impl(at::Tensor&, at::Tensor const&, bool) + 1944 (0x1625f6fe0 in libtorch_cpu.dylib)
frame #8: at::native::copy_(at::Tensor&, at::Tensor const&, bool) + 100 (0x1625f6788 in libtorch_cpu.dylib)
frame #9: c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor& (c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool), &(torch::ADInplaceOrView::copy_(c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool))>, at::Tensor&, c10::guts::typelist::typelist<c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool> >, at::Tensor& (c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool) + 76 (0x1665521a8 in libtorch_cpu.dylib)
frame #10: c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor& (c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool), &(torch::autograd::VariableType::(anonymous namespace)::copy_(c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool))>, at::Tensor&, c10::guts::typelist::typelist<c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool> >, at::Tensor& (c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool) + 772 (0x16654f880 in libtorch_cpu.dylib)
frame #11: at::_ops::copy_::call(at::Tensor&, at::Tensor const&, bool) + 288 (0x16332d0f4 in libtorch_cpu.dylib)
frame #12: torch::autograd::generated::details::fft_r2c_backward(at::Tensor const&, c10::ArrayRef<long long>, long long, bool, c10::SymInt) + 788 (0x1664f7c74 in libtorch_cpu.dylib)
frame #13: torch::autograd::generated::FftR2CBackward0::apply(std::__1::vector<at::Tensor, std::__1::allocator<at::Tensor> >&&) + 312 (0x164b31144 in libtorch_cpu.dylib)
frame #14: torch::autograd::Node::operator()(std::__1::vector<at::Tensor, std::__1::allocator<at::Tensor> >&&) + 120 (0x165a92008 in libtorch_cpu.dylib)
frame #15: torch::autograd::Engine::evaluate_function(std::__1::shared_ptr<torch::autograd::GraphTask>&, torch::autograd::Node*, torch::autograd::InputBuffer&, std::__1::shared_ptr<torch::autograd::ReadyQueue> const&) + 2932 (0x165a88df4 in libtorch_cpu.dylib)
frame #16: torch::autograd::Engine::thread_main(std::__1::shared_ptr<torch::autograd::GraphTask> const&) + 640 (0x165a87c98 in libtorch_cpu.dylib)
frame #17: torch::autograd::Engine::thread_init(int, std::__1::shared_ptr<torch::autograd::ReadyQueue> const&, bool) + 336 (0x165a8697c in libtorch_cpu.dylib)
frame #18: torch::autograd::python::PythonEngine::thread_init(int, std::__1::shared_ptr<torch::autograd::ReadyQueue> const&, bool) + 112 (0x1059f5898 in libtorch_python.dylib)
frame #19: void* std::__1::__thread_proxy<std::__1::tuple<std::__1::unique_ptr<std::__1::__thread_struct, std::__1::default_delete<std::__1::__thread_struct> >, void (torch::autograd::Engine::*)(int, std::__1::shared_ptr<torch::autograd::ReadyQueue> const&, bool), torch::autograd::Engine*, signed char, std::__1::shared_ptr<torch::autograd::ReadyQueue>, bool> >(void*) + 76 (0x165a95168 in libtorch_cpu.dylib)
frame #20: _pthread_start + 148 (0x1bf01426c in libsystem_pthread.dylib)
frame #21: thread_start + 8 (0x1bf00f08c in libsystem_pthread.dylib)

zsh: abort      svc train
```





Discussion:
I believe there was a feature to do that automatically (in lightning?)
@34j 
Oh, okay with AMP. You're right. Why do you think this is bugging out?
Maybe the issue happens when executing fft_r2c_backward 

Perhaps, it's an issue to be raised with PyTorch
It seems that environmental variable PYTORCH_ENABLE_MPS_FALLBACK should be set to 1, please try that
https://stackoverflow.com/a/72416727
Thank you, @34j 
```
$ conda env config vars list

PYTORCH_ENABLE_MPS_FALLBACK = 1
```

But still when I run `svc train -t`

I'm still getting:
```
Training: 0it [00:00, ?it/s]           INFO     [15:24:11] Setting current epoch to 0           train.py:198
           INFO     [15:24:11] Setting total batch idx to 0         train.py:213
           INFO     [15:24:11] Setting global step to 0             train.py:203
Epoch 0:   0%|                                           | 0/10 [00:00<?, ?it/s]libc++abi: terminating with uncaught exception of type c10::Error: Unsupported type byte size: ComplexFloat
Exception raised from getGatherScatterScalarType at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/mps/operations/View.mm:758 (most recent call first):
frame #0: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > const&) + 92 (0x1054e92b8 in libc10.dylib)
frame #1: at::native::mps::getGatherScatterScalarType(at::Tensor const&) + 304 (0x15e92266c in libtorch_cpu.dylib)
frame #2: invocation function for block in at::native::mps::gatherViewTensor(at::Tensor const&, at::Tensor&) + 128 (0x15e9241bc in libtorch_cpu.dylib)
frame #3: _dispatch_client_callout + 20 (0x1bee501b4 in libdispatch.dylib)
frame #4: _dispatch_lane_barrier_sync_invoke_and_complete + 56 (0x1bee5f414 in libdispatch.dylib)
frame #5: at::native::mps::gatherViewTensor(at::Tensor const&, at::Tensor&) + 888 (0x15e922d54 in libtorch_cpu.dylib)
frame #6: at::native::mps::mps_copy_(at::Tensor&, at::Tensor const&, bool) + 3096 (0x15e87a47c in libtorch_cpu.dylib)
frame #7: at::native::copy_impl(at::Tensor&, at::Tensor const&, bool) + 1944 (0x15a5f6fe0 in libtorch_cpu.dylib)
frame #8: at::native::copy_(at::Tensor&, at::Tensor const&, bool) + 100 (0x15a5f6788 in libtorch_cpu.dylib)
frame #9: c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor& (c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool), &(torch::ADInplaceOrView::copy_(c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool))>, at::Tensor&, c10::guts::typelist::typelist<c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool> >, at::Tensor& (c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool) + 76 (0x15e5521a8 in libtorch_cpu.dylib)
frame #10: c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor& (c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool), &(torch::autograd::VariableType::(anonymous namespace)::copy_(c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool))>, at::Tensor&, c10::guts::typelist::typelist<c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool> >, at::Tensor& (c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor&, at::Tensor const&, bool) + 772 (0x15e54f880 in libtorch_cpu.dylib)
frame #11: at::_ops::copy_::call(at::Tensor&, at::Tensor const&, bool) + 288 (0x15b32d0f4 in libtorch_cpu.dylib)
frame #12: torch::autograd::generated::details::fft_r2c_backward(at::Tensor const&, c10::ArrayRef<long long>, long long, bool, c10::SymInt) + 788 (0x15e4f7c74 in libtorch_cpu.dylib)
frame #13: torch::autograd::generated::FftR2CBackward0::apply(std::__1::vector<at::Tensor, std::__1::allocator<at::Tensor> >&&) + 312 (0x15cb31144 in libtorch_cpu.dylib)
frame #14: torch::autograd::Node::operator()(std::__1::vector<at::Tensor, std::__1::allocator<at::Tensor> >&&) + 120 (0x15da92008 in libtorch_cpu.dylib)
frame #15: torch::autograd::Engine::evaluate_function(std::__1::shared_ptr<torch::autograd::GraphTask>&, torch::autograd::Node*, torch::autograd::InputBuffer&, std::__1::shared_ptr<torch::autograd::ReadyQueue> const&) + 2932 (0x15da88df4 in libtorch_cpu.dylib)
frame #16: torch::autograd::Engine::thread_main(std::__1::shared_ptr<torch::autograd::GraphTask> const&) + 640 (0x15da87c98 in libtorch_cpu.dylib)
frame #17: torch::autograd::Engine::thread_init(int, std::__1::shared_ptr<torch::autograd::ReadyQueue> const&, bool) + 336 (0x15da8697c in libtorch_cpu.dylib)
frame #18: torch::autograd::python::PythonEngine::thread_init(int, std::__1::shared_ptr<torch::autograd::ReadyQueue> const&, bool) + 112 (0x106f51898 in libtorch_python.dylib)
frame #19: void* std::__1::__thread_proxy<std::__1::tuple<std::__1::unique_ptr<std::__1::__thread_struct, std::__1::default_delete<std::__1::__thread_struct> >, void (torch::autograd::Engine::*)(int, std::__1::shared_ptr<torch::autograd::ReadyQueue> const&, bool), torch::autograd::Engine*, signed char, std::__1::shared_ptr<torch::autograd::ReadyQueue>, bool> >(void*) + 76 (0x15da95168 in libtorch_cpu.dylib)
frame #20: _pthread_start + 148 (0x1bf01426c in libsystem_pthread.dylib)
frame #21: thread_start + 8 (0x1bf00f08c in libsystem_pthread.dylib)

zsh: abort      svc train -t
```


Try replacing here with if True: to patch stft
        if isinstance(self.trainer.accelerator, TPUAccelerator):
https://github.com/34j/so-vits-svc-fork/blob/main/src/so_vits_svc_fork/train.py#L179-L179
I did that and now my code looks like this: 

```
 # check if using tpu
        if True: to patch stft
        if isinstance(self.trainer.accelerator, TPUAccelerator):
            # patch torch.stft to use cpu
            LOG.warning("Using TPU. Patching torch.stft to use cpu.")
```
I get an error saying "TabError: inconsistent use of tabs and spaces in indentation"

I thought it wasnt possible to train on a mac?



I mean
```python
        if True:
            # patch torch.stft to use cpu
            LOG.warning("Using TPU. Patching torch.stft to use cpu.")
```
@34j 

It looks like that worked!

It's running thru the epochs. I will dedicate tomorrow to training a model and I'll let you know if I was successful.

Thank you so much for your help 🙏

Resolving patch:
diff --git a/src/so_vits_svc_fork/train.py b/src/so_vits_svc_fork/train.py
--- a/src/so_vits_svc_fork/train.py
+++ b/src/so_vits_svc_fork/train.py
@@ -8,7 +8,7 @@
 
 import lightning.pytorch as pl
 import torch
-from lightning.pytorch.accelerators import TPUAccelerator
+from lightning.pytorch.accelerators import MPSAccelerator, TPUAccelerator
 from lightning.pytorch.loggers import TensorBoardLogger
 from lightning.pytorch.tuner import Tuner
 from torch.cuda.amp import autocast
@@ -175,10 +175,10 @@ def on_train_start(self) -> None:
             global_step = total_batch_idx * self.optimizers_count
             self.set_global_step(global_step)
 
-        # check if using tpu
-        if isinstance(self.trainer.accelerator, TPUAccelerator):
+        # check if using tpu or mps
+        if isinstance(self.trainer.accelerator, (TPUAccelerator, MPSAccelerator)):
             # patch torch.stft to use cpu
-            LOG.warning("Using TPU. Patching torch.stft to use cpu.")
+            LOG.warning("Using TPU/MPS. Patching torch.stft to use cpu.")
 
             def stft(
                 input: torch.Tensor,
diff --git a/src/so_vits_svc_fork/utils.py b/src/so_vits_svc_fork/utils.py
--- a/src/so_vits_svc_fork/utils.py
+++ b/src/so_vits_svc_fork/utils.py
@@ -13,6 +13,7 @@
 import numpy as np
 import requests
 import torch
+import torch.backends.mps
 from cm_time import timer
 from fairseq import checkpoint_utils
 from fairseq.models.hubert.hubert import HubertModel
@@ -28,6 +29,8 @@
 def get_optimal_device(index: int = 0) -> torch.device:
     if torch.cuda.is_available():
         return torch.device(f"cuda:{index % torch.cuda.device_count()}")
+    elif torch.backends.mps.is_available():
+        return torch.device("mps")
     else:
         try:
             import torch_xla.core.xla_model as xm  # noqa


Test patch: