🐛 [BUG] Compatibility with torch1.9, geometric1.7, e3nn 0.3.2 about nequip HOT 5 CLOSED

opt_einsum_fx/opt_einsum_fx/_opt_ein.py:119: RuntimeWarning: einsum einsum lacked shape information; not optimizing. Did you forget to run ShapeProp on this graph?
  warnings.warn(

tests/data/test_AtomicData.py::test_non_periodic_edge[float32]

====================================================== FAILURES ======================================================
__________________________________________ test_non_periodic_edge[float32] ___________________________________________

CH3CHO = (Atoms(symbols='OCHCH3', pbc=False), AtomicData(edge_index=[2, 18], pos=[7, 3], num_nodes=7, atomic_numbers=[7], cell=[3, 3], edge_cell_shift=[18, 3], pbc=[3]))

    def test_non_periodic_edge(CH3CHO):
        atoms, data = CH3CHO
        # check edges
        for edge in range(data.num_edges):
            real_displacement = (
                atoms.positions[data.edge_index[1, edge]]
                - atoms.positions[data.edge_index[0, edge]]
            )
>           assert torch.allclose(
                data.get_edge_vectors()[edge],
                torch.as_tensor(real_displacement, dtype=torch.get_default_dtype()),
            )

tests/data/test_AtomicData.py:30: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

data = AtomicData(edge_index=[2, 18], pos=[7, 3], num_nodes=7, atomic_numbers=[7], cell=[3, 3], edge_cell_shift=[18, 3], pbc=[3])

    def get_edge_vectors(data: Data) -> torch.Tensor:
>       data = AtomicDataDict.with_edge_vectors(AtomicData.to_AtomicDataDict(data))
E       RuntimeError: with_edge_vectors() Expected a value of type 'Dict[str, Tensor]' for argument 'data' but instead found type 'dict'.
E       Position: 0
E       Value: {'cell': tensor([[0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.]]), 'pbc': tensor([False, False, False]), 'atomic_numbers': tensor([8, 6, 1, 6, 1, 1, 1]), 'num_nodes': 7, 'pos': tensor([[ 1.2181,  0.3612,  0.0000],
E               [ 0.0000,  0.4641,  0.0000],
E               [-0.4772,  1.4653,  0.0000],
E               [-0.9481, -0.7001,  0.0000],
E               [-0.3859, -1.6342,  0.0000],
E               [-1.5963, -0.6525,  0.8809],
E               [-1.5963, -0.6525, -0.8809]]), 'edge_cell_shift': tensor([[0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.],
E               [0., 0., 0.]]), 'edge_index': tensor([[0, 1, 1, 1, 2, 3, 3, 3, 3, 4, 4, 4, 5, 5, 5, 6, 6, 6],
E               [1, 3, 0, 2, 1, 5, 6, 4, 1, 3, 5, 6, 6, 4, 3, 5, 3, 4]])}
E       Declaration: with_edge_vectors(Dict(str, Tensor) data, bool with_lengths=True) -> (Dict(str, Tensor))
E       Cast error details: Unable to cast Python instance to C++ type (compile in debug mode for details)

nequip/data/AtomicData.py:253: RuntimeError

====================================================== FAILURES ======================================================
____________________________________________ test_without_nodes[float32] _____________________________________________

CH3CHO = (Atoms(symbols='OCHCH3', pbc=False), AtomicData(edge_index=[2, 18], pos=[7, 3], num_nodes=7, atomic_numbers=[7], cell=[3, 3], edge_cell_shift=[18, 3], pbc=[3]))

    def test_without_nodes(CH3CHO):
        # Non-periodic
        atoms, data = CH3CHO
        which_nodes = [0, 5, 6]
>       new_data = data.without_nodes(which_nodes=which_nodes)

tests/data/test_AtomicData.py:59: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = AtomicData(edge_index=[2, 18], pos=[7, 3], num_nodes=7, atomic_numbers=[7], cell=[3, 3], edge_cell_shift=[18, 3], pbc=[3])
which_nodes = tensor([0, 5, 6])

    def without_nodes(self, which_nodes):
        """Return a copy of ``self`` with ``which_nodes`` removed.
        The returned object may share references to some underlying data tensors with ``self``.
        Args:
            which_nodes (index tensor or boolean mask)
        Returns:
            A new data object.
        """
        which_nodes = torch.as_tensor(which_nodes)
        if which_nodes.dtype == torch.bool:
            mask = ~which_nodes
        else:
            mask = torch.ones(self.num_nodes, dtype=torch.bool)
            mask[which_nodes] = False
        assert mask.shape == (self.num_nodes,)
        n_keeping = mask.sum()
    
        # Only keep edges where both from and to are kept
        edge_mask = mask[self.edge_index[0]] & mask[self.edge_index[1]]
        # Create an index mapping:
        new_index = torch.full((self.num_nodes,), -1, dtype=torch.long)
        new_index[mask] = torch.arange(n_keeping, dtype=torch.long)
    
        new_dict = {}
        for k in self.keys:
            if k == AtomicDataDict.EDGE_INDEX_KEY:
                new_dict[AtomicDataDict.EDGE_INDEX_KEY] = new_index[
                    self.edge_index[:, edge_mask]
                ]
            elif k == AtomicDataDict.EDGE_CELL_SHIFT_KEY:
                new_dict[AtomicDataDict.EDGE_CELL_SHIFT_KEY] = self.edge_cell_shift[
                    edge_mask
                ]
            elif k == AtomicDataDict.CELL_KEY:
                new_dict[k] = self[k]
            else:
>               if len(self[k]) == self.num_nodes:
E               TypeError: object of type 'int' has no len()

nequip/data/AtomicData.py:325: TypeError