I wish to build a neural network that is both amenable to autograd and includes complex-values for input and output. For the latter, I applied torch.cfloat to change the dtypes in the neural network. However, I’m struggling to fix this error which I get after running my code. This is a snippet:
import numpy as np
import torch
import torch.optim as optim
from torch.autograd import grad
dtype=torch.cfloat

class mySin(torch.nn.Module):
    @staticmethod
    def forward(input):
        return torch.sin(input)

class Net(torch.nn.Module):
    def __init__(self, D_hid=10):
        super(Net,self).__init__()

        self.actF = mySin()
        self.Ein    = torch.nn.Linear(1,1, dtype=torch.cfloat)
        self.Lin_1  = torch.nn.Linear(2, D_hid, dtype=torch.cfloat)
        self.Lin_2  = torch.nn.Linear(D_hid, D_hid, dtype=torch.cfloat)
        self.out    = torch.nn.Linear(D_hid, 1, dtype=torch.cfloat)

    def forward(self,t):
        In1 = self.Ein(torch.ones_like(t, dtype=torch.cfloat))
        L1 = self.Lin_1(torch.cat((t,In1),1))
        h1 = self.actF(L1)
        L2 = self.Lin_2(h1)
        h2 = self.actF(L2)
        out = self.out(h2)
        return out, In1

def dfx(x,f):
    # Calculate the derivative with auto-differention
    return grad([f], [x], grad_outputs=torch.ones(x.shape, dtype=dtype), create_graph=True)[0]

# differential equation residual 
def diffeq_residual(t,psi, E):
    psi_dx = dfx(t,psi)
    psi_ddx= dfx(t,psi_dx)
    f  = (psi_ddx)/2 + E*(psi)
    L  = (f.pow(2)).mean(); 
    return L 

# testing the neural network
net = Net()
input = torch.rand(4).reshape(-1,1)
input.requires_grad = True
nn, En = net(input)
Loss = diffeq_residual(input, nn, En)
print(Loss)