[Python] torch.load don't open model weights and variables

I have only model file doclayout_yolo.pt in same folder as code weights-script.py and code:

import torch
import torch.nn as nn

class ExampleModel(nn.Module):
def __init__(self):
super().__init__()
self.flatten = nn.Flatten()
self.linear_relu_stack = nn.Sequential(
nn.Linear(28*28, 512),
nn.ReLU(),
nn.Linear(512, 10))

def forward(self, x):
x = self.flatten(x)
logits = self.linear_relu_stack(x)
return logits

model = ExampleModel()
model.load_state_dict(torch.load('doclayout_yolo.pt', weights_only=False))
print(model)


When I try to run I get error with text:

ModuleNotFoundError No module named 'doclayout_yolo' File "C:\Users\just llava\weights-script.py", line 19, in <module> model.load_state_dict(torch.load('doclayout_yolo.pt', weights_only=False)) ModuleNotFoundError: No module named 'doclayout_yolo'

When I change code to

model.load_state_dict(torch.load('doclayout_yolo.pt', weights_only=True))


I get error with text:

UnpicklingError Weights only load failed. This file can still be loaded, to do so you have two options, [1mdo those steps only if you trust the source of the checkpoint[0m. (1) In PyTorch 2.6, we changed the default value of the weights_only argument in torch.load from False to True. Re-running torch.load with weights_only set to False will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source. (2) Alternatively, to load with weights_only=True please check the recommended steps in the following error message. WeightsUnpickler error: Unsupported global: GLOBAL doclayout_yolo.nn.tasks.YOLOv10DetectionModel was not an allowed global by default. Please usetorch.serialization.add_safe_globals([doclayout_yolo.nn.tasks.YOLOv10DetectionModel]) or the torch.serialization.safe_globals([doclayout_yolo.nn.tasks.YOLOv10DetectionModel]) context manager to allowlist this global if you trust this class/function.

Check the documentation of torch.load to learn more about types accepted by default with weights_only https://pytorch.org/docs/stable/generated/torch.load.html. _pickle.UnpicklingError: Unsupported global: GLOBAL doclayout_yolo.nn.tasks.YOLOv10DetectionModel was not an allowed global by default. Please use torch.serialization.add_safe_globals([doclayout_yolo.nn.tasks.YOLOv10DetectionModel]) or the torch.serialization.safe_globals([doclayout_yolo.nn.tasks.YOLOv10DetectionModel]) context manager to allowlist this global if you trust this class/function.

During handling of the above exception, another exception occurred:

File "C:\Users\just llava\weights-script.py", line 19, in <module> model.load_state_dict(torch.load('doclayout_yolo.pt', weights_only=True)) _pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, [1mdo those steps only if you trust the source of the checkpoint[0m. (1) In PyTorch 2.6, we changed the default value of the weights_only argument in torch.load from False to True. Re-running torch.load with weights_only set to False will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source. (2) Alternatively, to load with weights_only=True please check the recommended steps in the following error message. WeightsUnpickler error: Unsupported global: GLOBAL doclayout_yolo.nn.tasks.YOLOv10DetectionModel was not an allowed global by default. Please use torch.serialization.add_safe_globals([doclayout_yolo.nn.tasks.YOLOv10DetectionModel]) or the torch.serialization.safe_globals([doclayout_yolo.nn.tasks.YOLOv10DetectionModel]) context manager to allowlist this global if you trust this class/function. Check the documentation of torch.load to learn more about types accepted by default with weights_only https://pytorch.org/docs/stable/generated/torch.load.html.

How to solve problem?

Continue reading...