class InvertibleResidualInplaceFunction(torch.autograd.function.Function):
		bug = False

		@staticmethod
		def forward(ctx, nonlinearity, x, *residual):
			ctx.nonlinearity = nonlinearity
			assert ctx.nonlinearity and ctx.nonlinearity[0] in ['leaky_relu']
			
			y = x.data
			for r in residual:
				y += r
			y = getattr(F, ctx.nonlinearity[0])(y, *ctx.nonlinearity[1:], inplace = True)
			ctx.num_residual = len(residual)
			
			if ResidualActivation.InvertibleResidualInplaceFunction.bug:
				residual = []
			
			#ctx.mark_dirty(x)
			ctx.save_for_backward(x, *residual)
			return x

		@staticmethod
		def backward(ctx, grad_output):
			x, *residual = ctx.saved_tensors
			y = x.data
			mask = torch.ones_like(grad_output).masked_fill_(x < 0, ctx.nonlinearity[1])
			grad_output *= mask
			y /= mask
			for r in residual:
				y -= r
			return (None, ) + (grad_output,) * (1 + ctx.num_residual)

class InplaceBatchNorm1d(nn.BatchNorm1d):
	def forward(self, input):
		return InplaceBatchNorm1d.Function.apply(input, self.weight, self.bias, self.running_mean, self.running_var, self.eps, self.momentum, self.training)

	class Function(torch.autograd.function.Function):
		@staticmethod
		def forward(ctx, input, weight, bias, running_mean, running_var, eps, momentum, training):
			mean, var = torch.batch_norm_update_stats(input, running_mean, running_var, momentum) if training else (running_mean, running_var)
			invstd = (var + eps).rsqrt_()
			output = torch.batch_norm_elemt(input, weight, bias, mean, invstd, 0, out = input)
			ctx.training = training
			ctx.save_for_backward(input, weight, bias, mean, invstd)
			ctx.mark_dirty(input)
			return input

		@staticmethod
		def backward(ctx, grad_output):
			assert ctx.training
			saved_output, weight, bias, mean, invstd = ctx.saved_tensors
			saved_input = torch.batch_norm_elemt(saved_output, invstd.reciprocal(), mean, bias, weight.reciprocal(), 0, out = saved_output)
			sum_dy, sum_dy_xmu, grad_weight, grad_bias = torch.batch_norm_backward_reduce(grad_output, saved_input, mean, invstd, weight, *ctx.needs_input_grad[:3])
			divisor = saved_input.numel() // saved_input.size(1)
			mean_dy = sum_dy.div_(divisor)
			mean_dy_xmu = sum_dy_xmu.div_(divisor)
			grad_input = torch.batch_norm_backward_elemt(grad_output, saved_input, mean, invstd, weight, mean_dy, mean_dy_xmu)
			return grad_input, grad_weight, grad_bias, None, None, None, None, None

# later on
residual = [...]
x = ...
residual_inputs = [bn(conv(r)) for conv, bn, r in zip(self.conv_residual, self.bn_residual, residual)]
x = self.activation(bn(conv(x)), residual = residual_inputs)
# activation applies the activation function, and bn is instance of this inplace batchnorm