rixty
/
KoboldAI-Client
mirror of https://github.com/KoboldAI/KoboldAI-Client.git
1
1
Fork 0
Code Issues Projects Releases Wiki Activity

Silence PyTorch warning about transposing tensors with dimension != 2

This commit is contained in:
Gnome Ann 2022-03-17 15:16:56 -04:00
parent ef21ab9c91
commit c444260eac
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tpu_mtj_backend.py
View File

@ -988,7 +988,7 @@ def load_model(path: str, driver_version="tpu_driver0.1_dev20210607", hf_checkpo
tensor /= params["cores_per_replica"] tensor /= params["cores_per_replica"]
if "vocab_pad" in transforms: if "vocab_pad" in transforms:
tensor = torch.nn.functional.pad(tensor, (0, 0, 0, params["n_vocab_padding"])) tensor = torch.nn.functional.pad(tensor, (0, 0, 0, params["n_vocab_padding"]))
if "no_transpose" not in transforms: if "no_transpose" not in transforms and tensor.ndim == 2:
tensor = tensor.T tensor = tensor.T
tensor.unsqueeze_(0) tensor.unsqueeze_(0)
if tensor.dtype is torch.float16 or tensor.dtype is torch.float32: if tensor.dtype is torch.float16 or tensor.dtype is torch.float32: