Hi,
Make sure to read the note on implementing custom Functions here 6.
In this case, your backward needs to return as many things as the forward had inputs. In your case, no gradient is needed for batch_size as I assume it’s a python number so you should return None for it in the backward.
Also note that you can avoid loops like this:
        for kk in range(batch_size):
            part1[kk] = part1[kk]*(alpha)

By using broadcasting:
        part1 = part1 * alpha

Also this won’t change part1 inplace (just change the Tensor the python object points to). So you won’t need the clone.