Solved by takafusui in post #2 
              
              
                Sorry this is my bad. 
In the above snippet, type(x) is torch.tensor and I overwrite the tensor object by torch.tensor(x[0]...) where I lost the grad_fn information, I guess. 
When I modify 
state_plus = torch.stack([x[0], aplus_val])[None, :]

everything works well.