After the execution of code mentioned below,The output shape of c1 variable after using torch.count_nonzero is torch.Size([]).

1.) When “vb.backward(torch.ones(va.size()))” is called, it gives following error as c1 shape is torch.Size([])

**RuntimeError: Mismatch in shape: grad_output[0] has a shape of torch.Size([1, 1, 256, 256]) and output[0] has a shape of torch.Size([]).**

2.) If I replace “c1=torch.count_nonzero(input)” with “c1 = torch.tensor([1.1], dtype=torch.float)” , then it gives the following error.

**RuntimeError: Mismatch in shape: grad_output[0] has a shape of torch.Size([1, 1, 256, 256]) and output[0] has a shape of torch.Size([1]).**

Code is as below.

class fun(Function):

@staticmethod

def forward(cxt,input):

cxt.save_for_backward(input)

c1=torch.count_nonzero(input)

print(c1.shape)

return c1

@staticmethod

def backward(cxt, grad_output):

input, = cxt.saved_tensors

grad_input = grad_output.clone()

grad_input[input < 0] = 0

return grad_input

fun1 = fun.apply

if **name** == “**main**”:

from torch.autograd import Variable

a=torch.rand(1,1,256,256)

va = Variable(a, requires_grad=True)

vb = fun1(va)

vb.backward(torch.ones(va.size()))