IndexError                                Traceback (most recent call last)
Cell In[5], line 1
----> 1 prediction = chgnet.predict_structure(structure)
      3 for key, unit in [
      4     ("energy", "eV/atom"),
      5     ("forces", "eV/A"),
      6     ("stress", "GPa"),
      7     ("magmom", "mu_B"),
      8 ]:
      9     print(f"CHGNet-predicted {key} ({unit}):\n{prediction[key[0]]}\n")

File ~/.conda/envs/LiIonMl/lib/python3.10/site-packages/chgnet/model/model.py:551, in CHGNet.predict_structure(self, structure, task, return_site_energies, return_atom_feas, return_crystal_feas, batch_size)
    548 structures = [structure] if isinstance(structure, Structure) else structure
    550 graphs = [self.graph_converter(struct) for struct in structures]
--> 551 return self.predict_graph(
    552     graphs,
    553     task=task,
    554     return_site_energies=return_site_energies,
    555     return_atom_feas=return_atom_feas,
    556     return_crystal_feas=return_crystal_feas,
    557     batch_size=batch_size,
    558 )

File ~/.conda/envs/LiIonMl/lib/python3.10/site-packages/chgnet/model/model.py:603, in CHGNet.predict_graph(self, graph, task, return_site_energies, return_atom_feas, return_crystal_feas, batch_size)
    601 n_steps = math.ceil(len(graphs) / batch_size)
    602 for step in range(n_steps):
--> 603     prediction = self.forward(
    604         [
    605             g.to(model_device)
    606             for g in graphs[batch_size * step : batch_size * (step + 1)]
    607         ],
    608         task=task,
    609         return_site_energies=return_site_energies,
    610         return_atom_feas=return_atom_feas,
    611         return_crystal_feas=return_crystal_feas,
    612     )
    613     for key in {
    614         "e",
    615         "f",
   (...)
    620         "crystal_fea",
    621     } & {*prediction}:
    622         for idx, tensor in enumerate(prediction[key]):

File ~/.conda/envs/LiIonMl/lib/python3.10/site-packages/chgnet/model/model.py:340, in CHGNet.forward(self, graphs, task, return_site_energies, return_atom_feas, return_crystal_feas)
    335 comp_energy = (
    336     0 if self.composition_model is None else self.composition_model(graphs)
    337 )
    339 # Make batched graph
--> 340 batched_graph = BatchedGraph.from_graphs(
    341     graphs,
    342     bond_basis_expansion=self.bond_basis_expansion,
    343     angle_basis_expansion=self.angle_basis_expansion,
    344     compute_stress="s" in task,
    345 )
    347 # Pass to model
    348 prediction = self._compute(
    349     batched_graph,
    350     compute_force="f" in task,
   (...)
    355     return_crystal_feas=return_crystal_feas,
    356 )

File ~/.conda/envs/LiIonMl/lib/python3.10/site-packages/chgnet/model/model.py:754, in BatchedGraph.from_graphs(cls, graphs, bond_basis_expansion, angle_basis_expansion, compute_stress)
    751 # Bonds
    752 atom_cart_coords = graph.atom_frac_coord @ lattice
    753 bond_basis_ag, bond_basis_bg, bond_vectors = bond_basis_expansion(
--> 754     center=atom_cart_coords[graph.atom_graph[:, 0]],
    755     neighbor=atom_cart_coords[graph.atom_graph[:, 1]],
    756     undirected2directed=graph.undirected2directed,
    757     image=graph.neighbor_image,
    758     lattice=lattice,
    759 )
    760 atom_positions.append(atom_cart_coords)
    761 bond_bases_ag.append(bond_basis_ag)

IndexError: tensors used as indices must be long, byte or bool tensors