fix bug in _einsum_flops_compute when using torch.einsum

[Necolizer]

Thank you for your excellent work!

This issue #40 gives a toy example of multi-head self-attention implemented using torch.einsum. Running that script gets an AttributeError raised in function _einsum_flops_compute(). To fix this, I made minimal modifications to the way of getting input_shapes list:

input_shapes = [o_element.shape for o in operands for o_element in (o if isinstance(o, list) else [o])]

If o is a list, we just collect the shapes of all its elements (o_elements), otherwise we get the shape of o.