import torch


def compute_vertex_normals(vertices, faces):
    """
    Compute each vertex normal
    The code is inspired by pytorch3d's _compute_vertex_normals in structures/meshes.py, v0.6.0
    This function is probably not differentiable though
    """
    pts = vertices[faces]
    cross_product = torch.cross(pts[:, 1] - pts[:, 0], pts[:, 2] - pts[:, 1], dim=1)
    vertex_normals = torch.zeros(vertices.shape, device=vertices.device)
    vertex_normals = vertex_normals.index_add(
        0, faces.flatten(), cross_product.repeat_interleave(3, dim=0)
    )
    vertex_normals = vertex_normals / vertex_normals.square().sum(1).sqrt()[..., None]
    return vertex_normals