How to manipulate the *args input of a function?

Question

I ran across this function that mostly does what I want, but I need to tweak it a bit.

The data I have looks like this:

import torch
import torch.nn as