Hi, I am new to Pytorch. I have written a custom basic network. The above training framework is taking 5 min of time, for one update_model function call / one iteration in the CPU, whereas taking approximately 7 min when using GPU. And is there a better way of accumulating loss values, as I want the loss from each layer and take the average of loss for an entire batch. So, is there a way to do batch training for a custom written forward function ?  And while training the model, the RAM is piling up, and I assuming it is because of the accumulating the loss from each batch. Am I assuming right ?
The below code is for one layer of the network
class DetNet(nn.Module):   
    def __init__(self):
        super(DetNet, self).__init__()
        self.W1 = torch.nn.Parameter(torch.randn(8*K,5*K, requires_grad =True, dtype=torch.float64))
        self.b1 = torch.nn.Parameter(torch.randn(8*K,1, requires_grad = True, dtype=torch.float64))
        self.W2 = torch.nn.Parameter(torch.randn(K, 8*K, requires_grad = True, dtype=torch.float64))
        self.b2 = torch.nn.Parameter(torch.randn(K, 1, requires_grad = True, dtype=torch.float64))
        self.W3 = torch.nn.Parameter(torch.randn(2*K, 8*K, requires_grad = True, dtype=torch.float64))
        self.b3 = torch.nn.Parameter(torch.randn(2*K, 1, requires_grad = True, dtype=torch.float64))
        self.t  = torch.nn.Parameter(torch.randn(1,1, requires_grad = True, dtype=torch.float64))

    def forward(self, x, v, y, H):
        M1 = torch.matmul(torch.transpose(H,0,1), y)
        M2 = torch.matmul(torch.transpose(H,0,1), torch.matmul(H,x))
        con = torch.cat((M1, x, M2, v))
        z = F.relu(torch.matmul(self.W1, con) + self.b1)
        y = torch.matmul(self.W2, z)+self.b2
        one_K = torch.ones([K,1])
        x_k = (F.relu(y+(one_K*self.t))/abs(self.t) - F.relu(y-(one_K*self.t))/abs(self.t)-one_K)
        v_k = torch.matmul(self.W3, z) + self.b3
        return (x_k, v_k)

The below code is for update of network for one epoch of the model.
def update_model(optimizer):
#     Number of samples for each iteration
    H = varying_channel()
    loss = torch.tensor(0).double()
    for samples in range(3500):
      # The below four lines are for generation of a random data sample 
        x_main = torch.DoubleTensor([[(2*round(np.random.rand())-1)] for cnt in range(K)])
        v = torch.zeros([2*K,1], dtype=torch.float64)
        y = received_signal(x_main, H)
        x_tilde = ZF_decoder(H, y)
        x = x_main
      # Passing through each layer for accumulation of loss.
        for cnt in range(3*K):
            foo = DetLayers[cnt]
            (x_loc,v_loc)= foo.forward(x, v, y, H)
            (x, v) = (x_loc, v_loc)
            curr_loss = (log(cnt+1)* (torch.sum((x_main-x)**2))/ torch.sum((x_main-x_tilde)**2))
            loss = loss + curr_loss
    loss = loss/3500
    optimizer.zero_grad()
    loss.backward() 
    optimizer.step()
    return(loss.item())

The below code is for creating a specified number of layers and running the model for specified iterations.
DetLayers = []
ParamLayers = list()
for cnt in range(3*K):
    curr_layer = DetNet()
    DetLayers.append(curr_layer)
    ParamLayers.extend(list(curr_layer.parameters()))

# creating the optimizer
optimizer = optim.Adam(ParamLayers, lr = 0.01)
# Number of iterations
for iterations in range(1):
    val = update_model(optimizer)
    print(val)

Thanks in advance.