PicklingError: Can't pickle <class 'discopy.tensor.Box[float64]'>: attribute lookup Box[float64] on discopy.tensor failed

Question

Im new to QNLP, Im currently using Lambek python library from Quantinuum. Im currently running this example code on Training: Classical case from their tutorials: https://cqcl.github.io/lambeq/tutorials/trainer_classical.html

Im getting the error:

---------------------------------------------------------------------------
PicklingError                             Traceback (most recent call last)
Cell In[12], line 1
----> 1 trainer.fit(train_dataset, val_dataset, eval_interval=1, log_interval=5)

File c:\Users\sagar\MySpace\Quantinuum_QNLPAssignment\qnlpLambek\lib\site-packages\lambeq\training\trainer.py:435, in Trainer.fit(self, train_dataset, val_dataset, log_interval, eval_interval, early_stopping_interval)
    428 for batch in tqdm(train_dataset,
    429                   desc='Batch',
    430                   total=train_dataset.batches_per_epoch,
    431                   disable=disable_tqdm,
    432                   leave=False,
    433                   position=2):
    434     step += 1
--> 435     t_loss = self._step_and_eval(
    436         batch,
    437         self.training_step,
    438         train_losses,
    439         self._train_eval_running,
    440         self.evaluate_on_train
    441     )
    443     self._to_tensorboard('train/step_loss', t_loss, step)
    444     status_bar.set_description(
    445         self._generate_stat_report(
    446             train_loss=t_loss,
   (...)
...
--> 134 diagrams = pickle.loads(pickle.dumps(diagrams))  # deepcopy, but faster
    135 for diagram in diagrams:
    136     for b in diagram.boxes:

PicklingError: Can't pickle <class 'discopy.tensor.Box[float64]'>: attribute lookup Box[float64] on discopy.tensor failed

in the Trainer.fit() method.

Attaching the code below for reference:

import torch
import os
from lambeq import BobcatParser, SpiderAnsatz, Dataset, PytorchModel, PytorchTrainer

# Constants
BATCH_SIZE = 30
EPOCHS = 30
LEARNING_RATE = 3e-2
SEED = 0

# Function to read data
def read_data(filename):
    labels, sentences = [], []
    with open(filename) as f:
        for line in f:
            t = float(line[0])
            labels.append([t, 1-t])
            sentences.append(line[1:].strip())
    return labels, sentences

# Read data
train_labels, train_data = read_data('./datasets/mc_train_data.txt')
val_labels, val_data = read_data('./datasets/mc_dev_data.txt')
test_labels, test_data = read_data('./datasets/mc_test_data.txt')

# Testing flag
TESTING = int(os.environ.get('TEST_NOTEBOOKS', '0'))
if TESTING:
    train_labels, train_data = train_labels[:2], train_data[:2]
    val_labels, val_data = val_labels[:2], val_data[:2]
    test_labels, test_data = test_labels[:2], test_data[:2]
    EPOCHS = 1

# Parse sentences
parser = BobcatParser(verbose='text')
train_diagrams = parser.sentences2diagrams(train_data)
val_diagrams = parser.sentences2diagrams(val_data)
test_diagrams = parser.sentences2diagrams(test_data)

# Define ansatz
ansatz = SpiderAnsatz({AtomicType.NOUN: Dim(2), AtomicType.SENTENCE: Dim(2)})

# Create circuits
train_circuits = [ansatz(diagram) for diagram in train_diagrams]
val_circuits =  [ansatz(diagram) for diagram in val_diagrams]
test_circuits = [ansatz(diagram) for diagram in test_diagrams]

# Create PytorchModel
all_circuits = train_circuits + val_circuits + test_circuits
model = PytorchModel.from_diagrams(all_circuits)

# Define accuracy function
sig = torch.sigmoid
def accuracy(y_hat, y):
    return torch.sum(torch.eq(torch.round(sig(y_hat)), y))/len(y)/2

# Define evaluation metrics
eval_metrics = {"acc": accuracy}

# Create datasets
train_dataset = Dataset(train_circuits, train_labels, batch_size=BATCH_SIZE)
val_dataset = Dataset(val_circuits, val_labels, shuffle=False)

# Create trainer
trainer = PytorchTrainer(
        model=model,
        loss_function=torch.nn.BCEWithLogitsLoss(),
        optimizer=torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE),
        epochs=EPOCHS,
        evaluate_functions=eval_metrics,
        evaluate_on_train=True,
        verbose='text',
        seed=SEED)

# Fit the model
trainer.fit(train_dataset, val_dataset, eval_interval=1, log_interval=5)