Add custom trainable parameters in PyTorch

CustomTrainingParams.py

import random
import torch
import torch.nn as nn
from torch.autograd import Variable
from torch import optim
import torch.nn.functional as F

seq_len = 10
n_class = 5
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.W = nn.Parameter(torch.randn(seq_len, n_class), requires_grad=True)
        
    def forward(self, x):
        x = torch.matmul(x, self.W)
        x = F.relu(x)
        x = F.log_softmax(x)
        return x

model = SimpleNet()
print('parameters----------')
for parameter in model.parameters():
    print(parameter)

optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.0001)
loss_fn = nn.NLLLoss()
for i in range(10000):
    in_data = Variable(torch.FloatTensor([[random.randint(0, n_class-1) for _ in range(seq_len)]]))
    target_data = Variable(torch.LongTensor([ int(in_data.data[0][0]) ]))
    
    optimizer.zero_grad()
    output = model(in_data)
    loss = loss_fn(output, target_data)
    loss.backward()
    optimizer.step()
    if i % 1000 == 0:
        print('loss:', loss.data[0])
    
# test
print('======Test======')
for i in range(10):
    in_data = Variable(torch.FloatTensor([[random.randint(0, n_class-1) for _ in range(seq_len)]]))
    target_data = int(in_data.data[0][0])
    out = model(in_data)
    _, ind = torch.max(out.data[0], 0)
    print(ind[0] == target_data, 'out', ind[0], 'label', target_data)

print('after parameters')
print(next(model.parameters()))