neonbjb/tortoise-tts
1
0
Fork 0
mirror of https://github.com/neonbjb/tortoise-tts.git synced 2026-01-06 08:40:24 +01:00
Code Issues Actions Packages Projects Releases Wiki Activity

import deepspeed only if use_deepspeed is True

Browse source
This commit is contained in:
manmay-nakhashi 2023-07-10 07:47:46 +05:30
parent 31a2e153ff
commit 82724cca54
1 changed files with 1 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
tortoise/models/autoregressive.py
View file

@ -3,7 +3,6 @@ import functools
import torch
import torch.nn as nn
import torch.nn.functional as F
import deepspeed
from transformers import GPT2Config, GPT2PreTrainedModel, LogitsProcessorList
from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
from transformers.utils.model_parallel_utils import get_device_map, assert_device_map
@ -362,6 +361,7 @@ class UnifiedVoice(nn.Module):
self.mel_head
)
if use_deepspeed:
import deepspeed
self.ds_engine = deepspeed.init_inference(model=self.inference_model,
mp_size=1,
replace_with_kernel_inject=True,