An example of passing mutliple dictionaries
  a = ({'weight': torch.ones(1, 1)}, {'buffer': torch.zeros(1)})  # two separate dictionaries
  mod = nn.Bar(1, 1)  # return self.weight @ x + self.buffer
  print(mod.weight)  # tensor(...)
  print(mod.buffer)  # tensor(...)
  x = torch.randn((1, 1))
  print(x)
  functional_call(mod, a, x)  # same as x
  print(mod.weight)  # same as before functional_call

And here is an example of applying the grad transform over the
   parameters of a model.
  import torch
  import torch.nn as nn
  from torch.func import functional_call, grad

  x = torch.randn(4, 3)
  t = torch.randn(4, 3)
  model = nn.Linear(3, 3)

  def compute_loss(params, x, t):
      y = functional_call(model, params, x)
      return nn.functional.mse_loss(y, t)

  grad_weights = grad(compute_loss)(dict(model.named_parameters()), x, t)

Note:
 If the user does not need grad tracking outside of grad
