hi
Please see below for a example:

from numpy import binary_repr
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as utils
import torch.nn.init as init


class ThreeLayerNet(torch.nn.Module):
def __init__(self,D_in,H,D_out):
    super(ThreeLayerNet,self).__init__()
    self.linear1=torch.nn.Linear(D_in,H)
    self.linear2=torch.nn.Linear(H,H)
    self.linear3=torch.nn.Linear(H,D_out)

def forward(self,x):
    h_relu_1=self.linear1(x).clamp(min=0)
    h_relu_2=self.linear2(h_relu_1).clamp(min=0)
    y_pred=self.linear3(h_relu_2)
    return y_pred

def weight_init(m):
if isinstance(m,nn.Linear):
    init.xavier_normal_(m.weight.data)
    init.normal_(m.bias.data)

if name == ‘main’:
device=torch.device("cpu")
print('a')
D_in,D_out=76,45
H=128
model=ThreeLayerNet(D_in,H,D_out)
print('b')
model.apply(weight_init)
print('c')
criterion=nn.CrossEntropyLoss()
optimizer=torch.optim.Adam(model.parameters(),lr=1e-3)

x=torch.tensor([[0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0.,0., 0., 0., 0., 0., 0., 0., 0., 0.,0.,1., 0., 0., 0., 0., 0., 0., 0.,0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,0., 0., 0., 0.]])
y=torch.tensor([2])
print('d')
y_pred=model(x)
print('e')
loss_train=criterion(y_pred,y)
print('start bp')
loss_train.backward(retain_graph=True)
print('done bp')
optimizer.step()
print('done step')


use this code I can only print to d, and then get stuck.