rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

Unban Tensor

Browse Source
This commit is contained in:
Henk
2023-05-15 22:21:22 +02:00
parent 56443bc7ea
commit c5100b4eab
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
modeling/lazy_loader.py
View File

@@ -209,6 +209,7 @@ class RestrictedUnpickler(pickle.Unpickler):
"ByteStorage",
"BoolStorage",
"BFloat16Storage",
"Tensor",
):
return getattr(torch, name)
elif module == "numpy.core.multiarray" and name == "scalar":
@@ -221,7 +222,7 @@ class RestrictedUnpickler(pickle.Unpickler):
# Forbid everything else.
qualified_name = name if module == "__builtin__" else f"{module}.{name}"
raise pickle.UnpicklingError(
f"`{qualified_name}` is forbidden; the model you are loading probably contains malicious code"
f"`{qualified_name}` is forbidden; the model you are loading probably contains malicious code. If you think this is incorrect ask the developer to unban the ability for {module} to execute {name}"
)
def load(self, *args, **kwargs):