Oops, seems like that delay was mostly caused by me not removing
torch.autograd.set_detect_anomaly(True) 

after a debug. Might be “mostly”, as this
class HalfReLU(nn.Module):
	def __init__(self):
		super(HalfReLU, self).__init__()

	def forward(self, x):
		half = x.shape[1]//2
		y = nn.functional.relu(x[:, :half], inplace=True)
		return torch.cat(( y, x[:, half:] ), dim=1)

still works on the order of 20x slower than plain ReLU, but I’ll post back later only if I’ll find it to have severe impact on what I’m doing.