SFInfector IndexError

[samarthswarup]

I am trying to run the cpu version exactly as described, and encountering the following error:

---- MAIN IMPORTS SUCCESSFUL -----
============================================================
state MA week 202036
Seed used for python random, numpy and torch is 1234
devices used: [device(type='cpu')]
Run:  0
0 ('25015',)
Num Agents:  160775
exception: did not work for MA week 202036: index -9223372036854775808 is out of bounds for dimension 0 with size 5

Traceback (most recent call last):
  File "/Users/biocomplexity/scripts/python/GradABM/main.py", line 85, in run_all_weeks
    counties_predicted, predictions, learned_params = train_predict(args)
                                                     ^^^^^^^^^^^^^^^^^^^
  File "/Users/biocomplexity/scripts/python/GradABM/train_abm.py", line 519, in train_predict
    counties_predicted, predictions, learned_params = runner(params, devices, verbose)
                                                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/biocomplexity/scripts/python/GradABM/train_abm.py", line 381, in runner
    predictions = forward_simulator(params,param_values,abm,training_num_steps,counties,devices)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/biocomplexity/scripts/python/GradABM/train_abm.py", line 282, in forward_simulator
    _, pred_t = abm[counties[c]].step(time_step, param_t[c].to(model_device))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/biocomplexity/scripts/python/GradABM/abm_model.py", line 492, in step
    lam_t = self.net(agents_data, learnable_params['r0_value'])
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/biocomplexity/miniforge3/envs/abm/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/biocomplexity/scripts/python/GradABM/abm_model.py", line 105, in forward
    return self.propagate(edge_index, x=x, edge_attr=edge_attr, t=t,
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/biocomplexity/miniforge3/envs/abm/lib/python3.11/site-packages/torch_geometric/nn/conv/message_passing.py", line 467, in propagate
    out = self.message(**msg_kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/biocomplexity/scripts/python/GradABM/abm_model.py", line 112, in message
    tmp = self.lam(x_i, x_j, edge_attr, t,
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/biocomplexity/scripts/python/GradABM/abm_model.py", line 52, in lam
    A_s_i = SFInfector[x_j[:,1].long()]
            ~~~~~~~~~~^^^^^^^^^^^^^^^^^
IndexError: index -9223372036854775808 is out of bounds for dimension 0 with size 5

Any idea what might be causing this? Thanks!