Merge pull request #14981 from wangshuai09/gpu_info_for_ascend
Add training support and change lspci for Ascend NPU
This commit is contained in:
commit
eee46a5094
|
@ -95,6 +95,7 @@ class HypernetworkModule(torch.nn.Module):
|
|||
zeros_(b)
|
||||
else:
|
||||
raise KeyError(f"Key {weight_init} is not defined as initialization!")
|
||||
devices.torch_npu_set_device()
|
||||
self.to(devices.device)
|
||||
|
||||
def fix_old_state_dict(self, state_dict):
|
||||
|
|
|
@ -230,7 +230,7 @@ class FrozenCLIPEmbedderWithCustomWordsBase(torch.nn.Module):
|
|||
for fixes in self.hijack.fixes:
|
||||
for _position, embedding in fixes:
|
||||
used_embeddings[embedding.name] = embedding
|
||||
|
||||
devices.torch_npu_set_device()
|
||||
z = self.process_tokens(tokens, multipliers)
|
||||
zs.append(z)
|
||||
|
||||
|
|
4
webui.sh
4
webui.sh
|
@ -158,9 +158,9 @@ then
|
|||
if echo "$gpu_info" | grep -q "AMD" && [[ -z "${TORCH_COMMAND}" ]]
|
||||
then
|
||||
export TORCH_COMMAND="pip install torch==2.0.1+rocm5.4.2 torchvision==0.15.2+rocm5.4.2 --index-url https://download.pytorch.org/whl/rocm5.4.2"
|
||||
elif echo "$gpu_info" | grep -q "Huawei" && [[ -z "${TORCH_COMMAND}" ]]
|
||||
elif npu-smi info 2>/dev/null
|
||||
then
|
||||
export TORCH_COMMAND="pip install torch==2.1.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu; pip install torch_npu"
|
||||
export TORCH_COMMAND="pip install torch==2.1.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu; pip install torch_npu==2.1.0"
|
||||
|
||||
fi
|
||||
fi
|
||||
|
|
Loading…
Reference in New Issue