AttributeError: Can't pickle local object 'BrainDataset.__init__..' when running the Neuro_Demo notebook

[gmeng92]

Hi. I have encountered the AttributionError when running the cell in the Neuro_Demo notebook in the example folder.
The cell is

n_classes = 31
n_epochs = 2
meshnet = MeshNet(n_channels=1, n_classes=n_classes)

logdir = "logs/meshnet_mindboggle"

optimizer = torch.optim.Adam(meshnet.parameters(), lr=0.02)


scheduler = OneCycleLR(optimizer, max_lr=.02,
                 epochs=n_epochs, steps_per_epoch=len(train_loaders['train']))

runner = CustomRunner()
runner.train(
    model=meshnet,
    optimizer=optimizer,
    scheduler=scheduler,
    loaders=train_loaders,
    num_epochs=n_epochs,
    logdir=logdir,
    callbacks=[CheckpointCallback(logdir=logdir)],
    verbose=True
)

The error message is detailed as

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[46], line 14
     10 scheduler = OneCycleLR(optimizer, max_lr=.02,
     11                  epochs=n_epochs, steps_per_epoch=len(train_loaders['train']))
     13 runner = CustomRunner()
---> 14 runner.train(
     15     model=meshnet,
     16     optimizer=optimizer,
     17     scheduler=scheduler,
     18     loaders=train_loaders,
     19     num_epochs=n_epochs,
     20     logdir=logdir,
     21     callbacks=[CheckpointCallback(logdir=logdir)],
     22     verbose=True
     23 )

File [~/anaconda3/lib/python3.10/site-packages/catalyst/runners/runner.py:320](https://file+.vscode-resource.vscode-cdn.net/Users/M295878/Projects/neuro/examples/~/anaconda3/lib/python3.10/site-packages/catalyst/runners/runner.py:320), in Runner.train(self, loaders, model, engine, trial, criterion, optimizer, scheduler, callbacks, loggers, seed, hparams, num_epochs, logdir, valid_loader, valid_metric, minimize_valid_metric, verbose, timeit, check, overfit, load_best_on_end, fp16, amp, apex, ddp)
    318 self._load_best_on_end = load_best_on_end
    319 # run
--> 320 self.run()

File [~/anaconda3/lib/python3.10/site-packages/catalyst/core/runner.py:678](https://file+.vscode-resource.vscode-cdn.net/Users/M295878/Projects/neuro/examples/~/anaconda3/lib/python3.10/site-packages/catalyst/core/runner.py:678), in IRunner.run(self)
    676 except (Exception, KeyboardInterrupt) as ex:
    677     self.exception = ex
...
     58 def dump(obj, file, protocol=None):
     59     '''Replacement for pickle.dump() using ForkingPickler.'''
---> 60     ForkingPickler(file, protocol).dump(obj)

AttributeError: "Can't pickle local object 'BrainDataset.__init__.<locals>.<lambda>''

Does anyone have idea on how to fix it? Thanks!