Problem is as the title says. I’m unsure of why this is happening, as when I’m printing from the backward() method of my loss function, it prints the correct tensor for the grad, but when printing the grad from the training loop, it prints tensor(1.).
In the following code, I use a lot of code from a library called Qiskit. The main value that needs to be passed to the grad function is in state_fidelity(), which returns a scalar.
Example of some output:
Training loss for batch 0: 0.008105693179143336
tensor(1.)
[[tensor(8.5487e-15), tensor(-2.8533e-14)]]
Training loss for batch 1: 0.568790329178134
tensor(1.)
[[tensor(0.), tensor(0.)]]

training_loop:
epochs = 10
    
    losses = []
    
    loss_fn = SingleQubitGaussianLoss.apply
    
    print(model)
    
    model.train(True)
    
    for epoch in range(epochs):
        print(f"\nEPOCH {epoch}")
        for i, data in enumerate(data_loader):
            inputs, labels = data
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = loss_fn(outputs, inputs, labels)
            loss.register_hook(lambda grad: print(grad))
            if len(losses) < points:
                losses.append(loss.item())
                print(f"Training loss for batch {i}: {loss.item()}")
            else:
                comp = (loss.item() - losses[i]) / losses[i] * 100
                losses[i] = loss.item()
                print(f"Training loss for batch {i}: {loss.item()}, {comp}% change")
            loss.backward()
            optimizer.step()

loss function:
class SingleQubitGaussianLoss(torch.autograd.Function):
    # TODO somehow make duration a parameter
    @staticmethod
    def forward(ctx, params, inits, labels):
        ctx.save_for_backward(params)
        eval_inits = []
        eval_labels = []
        infidelity_list = []
        for pred, init, label in zip(params, inits, labels):
            amp, ln_sigma = pred
            amp = amp.item()
            ln_sigma = ln_sigma.item()
            init_state = [torch.cos(init[0] / 2).item(), torch.exp(init[1] * 1.j).item() * \
                          torch.sin(init[0] / 2).item()]
            eval_inits.append(init_state)
            job = run_gaussian(duration=128,
                               amp=process_amp(amp),
                               sigma=process_ln_width(ln_sigma),
                               init_state=init_state)
            sv = job.result().get_statevector()
            actual_sv = [label[0].item(), label[1].item(), 0]
            eval_labels.append(actual_sv)
            infidelity_list.append(1 - state_fidelity(sv, actual_sv))
        ctx.eval_inits = eval_inits
        ctx.eval_labels = eval_labels
        ctx.infidelity_list = infidelity_list
        return torch.DoubleTensor([sum(infidelity_list) / len(infidelity_list)])[0]
            
    
    @staticmethod
    def backward(ctx, h=1e-3):
        params, = ctx.saved_tensors
        jacobian = []
        for pred, init_state, actual_sv, infidelity in \
            zip(params, ctx.eval_inits, ctx.eval_labels, ctx.infidelity_list):
            amp, ln_sigma = pred
            amp = amp.item()
            ln_sigma = ln_sigma.item()
            amp_job = run_gaussian(duration=128,
                                   amp=process_amp(amp + h),
                                   sigma=process_ln_width(ln_sigma),
                                   init_state=init_state)
            amp_sv = amp_job.result().get_statevector()
            amp_infid = 1 - state_fidelity(actual_sv, amp_sv)
            grad_amp = (amp_infid - infidelity) / h
            
            sigma_job = run_gaussian(duration=128,
                                     amp=process_amp(amp),
                                     sigma=process_ln_width(ln_sigma + h),
                                     init_state=init_state)
            sigma_sv = sigma_job.result().get_statevector()
            sigma_infid = 1 - state_fidelity(actual_sv, sigma_sv)
            grad_sigma = (sigma_infid - infidelity) / h
            jacobian.append([grad_amp, grad_sigma])
        print(jacobian)
        return torch.DoubleTensor(jacobian), None, None


def process_amp(real: float, imag: float = 0):
    # TODO implement imaginary amps
    return 1 / (1 + math.exp(-1*real))


def process_ln_width(ln_width: float):
    return ln_width


model:
class SingleQubitPulseModel(torch.nn.Module):
    def __init__(self):
        super(SingleQubitPulseModel, self).__init__()
        
        # TODO Find better activation function, ReLU keeps producing 0 because amp is small
        self.linear1 = torch.nn.Linear(2, 64)
        self.relu1 = torch.nn.ReLU()
        self.linear2 = torch.nn.Linear(64, 64)
        self.relu2 = torch.nn.ReLU()
        self.linear3 = torch.nn.Linear(64, 2)
        self.relu3 = torch.nn.ReLU()
    
    def forward(self, x):
        x = self.linear1(x)
        x = self.relu1(x)
        x = self.linear2(x)
        x = self.relu2(x)
        x = self.linear3(x)
        x = self.relu3(x)
        return x