two dimensional calculation seems not supported

[ttx002000]

Hi,
I am using the following config file

import sys

from absl import logging
from ferminet.utils import system
from ferminet import base_config
from ferminet import train


logging.get_absl_handler().python_handler.stream = sys.stdout
logging.set_verbosity(logging.INFO)


cfg = base_config.default()
cfg.system.electrons = (1,0)  # (alpha electrons, beta electrons)
cfg.system.ndim = 2
cfg.system.molecule = [system.Atom("X", (0., 0.))]

cfg.batch_size = 256
cfg.pretrain.iterations = 0
cfg.network.jastrow= 'none'
cfg.network.complex= True

train.train(cfg)

Supposedly this should run a free electron gas in two dimensions. But I am getting error output like

 File "/home/ttx002000/project/ferminet/ferminet/networks.py", line 469, in construct_input_features
    assert atoms.shape[1] == ndim

So is ndims really an adjustable parameter in ferminet?