TypeError: _nblist_pbc_cpu() missing 1 required positional argument: 'device'

[corinwagen]

I get an error when trying to run periodic calculations on my M3 Macbook Pro (CPU only, no CUDA).

    def nblists_torch_pbc(coord: Tensor, cell: Tensor, cutoff: float) -> Tuple[Tensor, Tensor, Tensor]:
        """Compute dense neighbor lists for periodic boundary conditions case.
        Coordinates must be in cartesian coordinates and be within the unit cell.
        Single crystal only, no support for batched coord or multiple unit cells.
        """
        assert coord.ndim == 2, "Expected 2D tensor for coord, got {coord.ndim}D"
        # non-PBC version
        device = coord.device

        reciprocal_cell = cell.inverse().t()
        inv_distances = reciprocal_cell.norm(2, -1)
        shifts = _calc_shifts(inv_distances, cutoff)
        d = torch.cdist(coord.unsqueeze(0), coord.unsqueeze(0) + (shifts @ cell).unsqueeze(1))
        conn_mat = ((d < cutoff) & (d > 0.1)).transpose(0, 1).contiguous()
        if device.type == "cuda" and _numba_cuda_available:
            _fn = _nblist_pbc_cuda
        else:
            _fn = _nblist_pbc_cpu
>       mat_idxj, mat_pad, mat_S = _fn(conn_mat, shifts)
E       TypeError: _nblist_pbc_cpu() missing 1 required positional argument: 'device'

Looking at the type signatures in nblist.py, _nblist_pbc_cuda takes two arguments (conn_mat and shifts) while _nblist_pbc_cpu takes three (conn_mat, shifts, and device). Is the implementation of nblists_torch_pbc just wrong?

Would something as simple as this work?

     if device.type == "cuda" and _numba_cuda_available:
            mat_idxj, mat_pad, mat_S = _nblist_pbc_cuda(conn_mat, shifts)
        else:
            mat_idxj, mat_pad, mat_S = _nblist_pbc_cpu(conn_mat, shifts, device)