AttributeError: Can't get attribute '__main__' on <module 'builtins' (built-in)> while loading the pruned model

[jenyaMi]

I use yolov8 pruning code, to be more precisely the updated version for the new version of ultralytics: https://github.com/chbw818/yolov8-prune-using-torch-pruning-/tree/main

I am struggling with loading the pruned model...

AttributeError: Can't get attribute 'main' on <module 'builtins' (built-in)>

AttributeError Traceback (most recent call last) in 16 args = parser.parse_args() 17 ---> 18 prune(args) in prune(args) 432 pruning_cfg['name'] = f"step_{i}_finetune" 433 pruning_cfg['batch'] = batch_size # restore batch size --> 434 model.train_v2(pruning=True, **pruning_cfg) 435 436 # post fine-tuning validation in train_v2(self, pruning, **kwargs) 337 338 self.trainer.hub_session = self.session # attach optional HUB session --> 339 self.trainer.train() 340 # Update model and cfg after training 341 if RANK in (-1, 0): /workspace/data/notebooks/belt_detection/ultralytics/ultralytics/engine/trainer.py in train(self) 202 203 else: --> 204 self._do_train(world_size) 205 206 def _setup_scheduler(self): /workspace/data/notebooks/belt_detection/ultralytics/ultralytics/engine/trainer.py in _do_train(self, world_size) 467 f"{(time.time() - self.train_time_start) / 3600:.3f} hours." 468 ) --> 469 self.final_eval() 470 if self.args.plots: 471 self.plot_metrics() in final_eval_v2(self) 272 for f in self.last, self.best: 273 if f.exists(): --> 274 strip_optimizer_v2(f) # strip optimizers 275 if f is self.best: 276 LOGGER.info(f'\nValidating {f}...') in strip_optimizer_v2(f, s) 284 Disabled half precision saving. originated from ultralytics/yolo/utils/torch_utils.py 285 """ --> 286 x = torch.load(f, map_location=torch.device('cpu')) 287 args = {**DEFAULT_CFG_DICT, **x['train_args']} # combine model args with default args, preferring model args 288 if x.get('ema'): /opt/conda/lib/python3.8/site-packages/torch/serialization.py in load(f, map_location, pickle_module, **pickle_load_args) 605 opened_file.seek(orig_position) 606 return torch.jit.load(opened_file) --> 607 return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args) 608 return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args) /opt/conda/lib/python3.8/site-packages/torch/serialization.py in _load(zip_file, map_location, pickle_module, pickle_file, **pickle_load_args) 880 unpickler = UnpicklerWrapper(data_file, **pickle_load_args) 881 unpickler.persistent_load = persistent_load --> 882 result = unpickler.load() 883 884 torch._utils._validate_loaded_sparse_tensors() /opt/conda/lib/python3.8/site-packages/torch/serialization.py in find_class(self, mod_name, name) 873 def find_class(self, mod_name, name): 874 mod_name = load_module_mapping.get(mod_name, mod_name) --> 875 return super().find_class(mod_name, name) 876 877 # Load the data (which may in turn use persistent_load to load tensors) AttributeError: Can't get attribute 'main' on <module 'builtins' (built-in)>

Maybe someone know what could be the problem here, I saw many people were able to run the code successfully. I use the last version of ultralytics 8.2.50