Hey!
Would mind sharing what exactly X is / where is it calculated? Or is it just arbitrary input data?
From what I can see, you’re basically taking your network and modifying the parameters and creating a copy, then comparing the difference in output via the original model and a copy of the model with the modified parameters? Is this correct?
When copying across you use copy_ which is an in-place copy. Perhaps, it’d be better to deepcopy the model rather than the model’s state_dict and modifying the parameters in that way all in one go rather than two separate loops. Something like this might work!
model_star = deepcopy(model)
for model_np, model_star_np  in  zip(model.named_parameters(), model_star.named_parameters()):
  name, param = model_np                #unpack both tuples (name and param) 
  name_star, param_star = model_star_np #for each model
 
  random_val = torch.randn(*param_star.shape, device=device, requires_grad=True)
  param_star.copy_( v + v.abs() * initial_std * random_val) #in-place copy across to model_star

#calculate both outputs here
output_theta = model(X)
output_theta_star = model_star(X)

#Calc. loss here
step_loss = robustness_loss(torch.nn.functional.softmax(output_theta, dim=1).log(), output_theta_star)

optim.zero_grad() #don't forget to zero your grads! (unless you just want the grad of loss w.r.t both model parameters?)
step_loss.backward()
optim.step()

When creating a custom loss function writing a standard python function should be fine. torch.autograd.Function is used when you want to create a custom function with a forward and backward method, for example, a custom activation function or module layer! For example a custom ReLU function example is in the docs here → PyTorch: Defining New autograd Functions — PyTorch Tutorials 1.7.0 documentation