Hi CEBRA Team, there is no attribute 'UnifiedDataset' from the DEMO of unified CEBRA encoders

[Lordawn]

When I followed demo code about unified CEBRA encoders cross different animals https://cebra.ai/docs/demo_notebooks/Demo_hippocampus_unified.html, it shown

AttributeError Traceback (most recent call last)
Cell In[5], line 30
27 num_neurons += dataset.neural.shape[1]
29 # Define the dataset
---> 30 dataset = cebra.data.UnifiedDataset(dataset for dataset in train_datasets)
32 # Set the masks for the dataset
33 # 'RandomNeuronMask' and 'RandomTimestepMask' are used to randomly mask neurons and timesteps
34 # during training, which helps the model to learn better representations.
35 dataset.set_masks({"RandomNeuronMask": (0.3, 0.9, 0.05),
36 "RandomTimestepMask": (0.3, 0.9, 0.05)})

AttributeError: module 'cebra.data' has no attribute 'UnifiedDataset'

it occured in both Windows and Linux system. I wondering if there is any update on replacing this 'UnifiedDataset' function on constructing a cross sessions/animals dataset?

[MMathisLab]

Which version did you install?

[Lordawn]

Following the instruction in https://cebra.ai/docs/installation.html, using pip install cebra in conda, i do not sure which version exactly

[mfanourakis]

I ran into the same problem. Looks like the Demos and API docs are all a little ahead in versions (>v0.6) compared to the latest version in pip (v0.5).
There are enough differences between these versions that the demos and API docs are almost useless for the torch API of CEBRA.

[MMathisLab]

You'll need to install with --pre tag / or from GitHub

[mfanourakis]

Thanks!
works as expected when installing the pre-release version from pip