Error in the `_grad_cache_forward_step` function when running suggested contrastive pretraining command

Hi, thank you very much for releasing the code and data for training text embedding models.

I followed the steps in README.md and ran the following suggested command for the contrastive pretraining step:

cd src/contrastors
torchrun --nproc-per-node=8 train.py --config=configs/train/contrastive_pretrain.yaml --dtype=bf16

However, I get the following error

File ".../contrastors/src/contrastors/trainers/clip.py", line 163, in _grad_cache_forward_step
    chunk_size=self.config.chunk_size,                                       
AttributeError: 'Config' object has no attribute 'chunk_size'

the source of it is: https://github.com/nomic-ai/contrastors/blob/70f7baf0dff4e0019776e8c98257a3ad3fa4d687/src/contrastors/trainers/clip.py#L151-L167

My guess is that one should change chunk_size=self.config.chunk_size to chunk_size=self.config.train_args.chunk_size.

However, after changing that I get another error in the same _grad_cache_forward_step function:

File ".../contrastors/src/contrastors/trainers/clip.py", line 158, in _grad_cache_forward_step                             
    loss = grad_cache_loss(                                                                                                                                            
TypeError: grad_cache_loss() got an unexpected keyword argument 'step'

which I'm guessing comes from running the forward_step function with step as an argument.

I see that the comment in the _grad_cache_forward_step function says that this function is not used? So is it that some code changes have not been pushed?

Many thanks, Kuan

nomic-ai / contrastors

Error in the `_grad_cache_forward_step` function when running suggested contrastive pretraining command #9