Simple JIT scripting fails with "Tried to access nonexistent attribute or method..."

I am new to JIT scripting and don’t know what this error means, and have not found this zero_grad error online nor any related errors whose solutions seem to make sense here. What am I doing wrong here?

The error is this (see gist for full error):

Tried to access nonexistent attribute or method 'zero_grad' of type 'Tensor (inferred)'.

The full repro is here:

I think the error is raised, because scripting assumes tensors as the input, if no annotations are used as described here.
I don’t know, if it’s possible to script the complete optimization, but you should be able to script the model at least.