Bugfix: parameters of both the model as well as the criterion (Adapti…

…ve Softmax) must be optimized

main.py

@@ -141,12 +141,11 @@ def model_load(fn):
  print('Using', splits)
  criterion = SplitCrossEntropyLoss(args.emsize, splits=splits, verbose=False)
 ###
-params = list(model.parameters()) + list(criterion.parameters())
 if args.cuda:
  model = model.cuda()
  criterion = criterion.cuda()
- params = list(model.parameters()) + list(criterion.parameters())
 ###
+params = list(model.parameters()) + list(criterion.parameters())
 total_params = sum(x.size()[0] * x.size()[1] if len(x.size()) > 1 else x.size()[0] for x in params if x.size())
 print('Args:', args)
 print('Model total parameters:', total_params)
@@ -206,7 +205,7 @@ def train():
  loss.backward()
 
  # `clip_grad_norm` helps prevent the exploding gradient problem in RNNs / LSTMs.
- if args.clip: torch.nn.utils.clip_grad_norm(model.parameters(), args.clip)
+ if args.clip: torch.nn.utils.clip_grad_norm(params, args.clip)
  optimizer.step()
 
  total_loss += raw_loss.data
@@ -232,10 +231,11 @@ def train():
 # At any point you can hit Ctrl + C to break out of training early.
 try:
  optimizer = None
+ # Ensure the optimizer is optimizing params, which includes both the model's weights as well as the criterion's weight (i.e. Adaptive Softmax)
  if args.optimizer == 'sgd':
- optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, weight_decay=args.wdecay)
+ optimizer = torch.optim.SGD(params, lr=args.lr, weight_decay=args.wdecay)
  if args.optimizer == 'adam':
- optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wdecay)
+ optimizer = torch.optim.Adam(params, lr=args.lr, weight_decay=args.wdecay)
  for epoch in range(1, args.epochs+1):
  epoch_start_time = time.time()
  train()