Tweak the code generation to properly mark self as mutable, closes #33.

gen/gen.ml

@@ -229,13 +229,18 @@ module Func = struct
           Printf.sprintf "%s: %s" (rust_name arg.arg_name) rust_arg_type )
       |> String.concat ~sep:", "
     in
+    let self_arg =
+      if String.is_suffix t.name ~suffix:"_"
+      then "&mut self"
+      else "&self"
+    in
     match List.partition_tf t.args ~f:self_tensor with
     | [self], args_list ->
-        (Some self.arg_name, Printf.sprintf "&self, %s" (to_string args_list))
+        (Some self.arg_name, Printf.sprintf "%s, %s" self_arg (to_string args_list))
     | _, _ -> (
       match List.partition_tf t.args ~f:input_tensor with
       | [self], args_list ->
-          (Some self.arg_name, Printf.sprintf "&self, %s" (to_string args_list))
+          (Some self.arg_name, Printf.sprintf "%s, %s" self_arg (to_string args_list))
       | _, _ -> (None, to_string t.args) )
 
   let rust_return_type t ~fallible =

src/wrappers/tensor.rs

@@ -129,7 +129,7 @@ impl Tensor {
 
     /// Zeroes the gradient tensor attached to this tensor if defined.
     pub fn zero_grad(&mut self) {
-        let grad = self.grad();
+        let mut grad = self.grad();
         if grad.defined() {
             let _ = grad.detach_().zero_();
         }