FilippNikitin

**Description** The line: ```python edge_types[edge_types == 4] = 1.5 ``` fails because edge_types is a long tensor, truncating 1.5 to 1. This replaces 4 with 1 instead of 1.5, affecting...