_pickle.PicklingError: Невозможно выбрать <class '__main __. X'>: поиск атрибута X на __main__ не выполнен

Question

Pickle успешно работает в ноутбуке Jupiter, но когда я делаю это в .py файле, Pickle завершается неудачей.

Я не понимаю, почему?

Это просто простой класс

class X:
def __init__(self, ref):
    import torch
    self.ref = torch.tensor(ref)
    self.result = torch.full(ref.shape, True, dtype=torch.uint8)
    self.index = None

def partial_fit(self, X):
    assert len(X) > 0
    import numpy
    blablabla

def transform(self, X):
    return X[:, self.index]

Версия и пакет Python выглядят одинаково