Firstly, Thanks for your great work!
I am using the counterfactual module computes counterfactual examples, When starting runing, it raises the following error:

:/data/wuyunxiao/ECINN-main$ CUDA_VISIBLE_DEVICES=2 python main.py counterfactual -c config.ini
['config.ini']

Using config: config.ini
Files already downloaded and verified
Files already downloaded and verified

# # # # # # # # # # # # # # # # # # # # # # # # # # # #

Computing counterfactuals

# # # # # # # # # # # # # # # # # # # # # # # # # # # #

Computing mean embeddings
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 389/389 [00:23<00:00, 16.64it/s]
None
Traceback (most recent call last):
File "main.py", line 78, in
main()
File "main.py", line 74, in main
args.module.fn(args, cfg, model, data)
File "/data/wuyunxiao/ECINN-main/counterfactual/init.py", line 14, in fn
compute_counterfactuals(args, cfg, model, data)
File "/home/wuyunxiao/test/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
return func(*args, **kwargs)
File "/data/wuyunxiao/ECINN-main/counterfactual/compute_and_store.py", line 50, in compute_counterfactuals
means = mean_embeddings(cfg, data, model)[1:] # Skip total mean
File "/home/wuyunxiao/test/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
return func(*args, **kwargs)
File "/data/wuyunxiao/ECINN-main/counterfactual/counterfactuals.py", line 53, in mean_embeddings
means = torch.cat([m.mean.unsqueeze(0) for m in [overall_mean] + means], 0)
File "/data/wuyunxiao/ECINN-main/counterfactual/counterfactuals.py", line 53, in
means = torch.cat([m.mean.unsqueeze(0) for m in [overall_mean] + means], 0)
AttributeError: 'NoneType' object has no attribute 'unsqueeze'

Could you give me some advice on how to get rid of it?