Function squared_relu

gpu/model.py:165–166

(x: torch.Tensor)

@torch.compile
def squared_relu(x: torch.Tensor) -> torch.Tensor:
    return F.relu(x) ** 2

class FeedForward(nn.Module):
    def __init__(