Skorch training object from scratch

I'm trying to use skorch class to execut GridSearch on a classifier. I tried running with the vanilla NeuralNetClassifier object, but I haven't found a way to pass the Adam optimizer only the trainable weights (I'm using pre-trained embeddings and I would like to keep them frozen). It's doable if a module is initialized, and then pass those weights with the optimizer__params option, but module needs an uninitialized model. Is there a way around this?

net = NeuralNetClassifier(module=RNN, module__vocab_size=vocab_size, module__hidden_size=hidden_size,
                          module__embedding_dim=embedding_dim, module__pad_id=pad_id,
                          module__dataset=ClaimsDataset, lr=lr, criterion=nn.CrossEntropyLoss,
                          optimizer=torch.optim.Adam, optimizer__weight_decay=35e-3, device='cuda',
                          max_epochs=nb_epochs, warm_start=True)

The code above works. However, with the batch_size set at 64, I've got to run the model for the specified number of epochs on every batch! Which is not the behavior I'm seeking. I'd be grateful if someone could suggest a nicer way to do this.

My other issue is with subclassing skorch.NeuralNet. I run into a similar issue: figuring out a way to pass only the trainable weights to Adam optimizer. The code below is what I've got so far.

class Train(skorch.NeuralNet):
def __init__(self, module, lr, norm, *args, **kwargs):
    self.module = module
    self.lr = lr
    self.norm = norm
    self.params = [p for p in self.module.parameters(self) if p.requires_grad]
    super(Train, self).__init__(*args, **kwargs)

def initialize_optimizer(self):
    self.optimizer = torch.optim.Adam(params=self.params, lr=self.lr, weight_decay=35e-3, amsgrad=True)

def train_step(self, Xi, yi, **fit_params):
    self.module.train()

    self.optimizer.zero_grad()
    yi = variable(yi)

    output = self.module(Xi)

    loss = self.criterion(output, yi)
    loss.backward()

    nn.utils.clip_grad_norm_(self.params, max_norm=self.norm)
    self.optimizer.step()

def score(self, y_t, y_p):
    return accuracy_score(y_t, y_p)

Initializing the class gives the error:

Traceback (most recent call last):
File "/snap/pycharm-community/74/helpers/pydev/pydevd.py", line 1664, in <module>
    main()
File "/snap/pycharm-community/74/helpers/pydev/pydevd.py", line 1658, in main
globals = debugger.run(setup['file'], None, None, is_module)
File "/snap/pycharm-community/74/helpers/pydev/pydevd.py", line 1068, in run
pydev_imports.execfile(file, globals, locals)  # execute the script
File "/snap/pycharm-community/74/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
File "/home/l/Documents/Bsrc/cv.py", line 115, in <module>
    main()
File "/home/l/B/src/cv.py", line 86, in main
    trainer = Train(module=RNN, criterion=nn.CrossEntropyLoss, lr=lr, norm=max_norm)
File "/home/l/B/src/cv.py", line 22, in __init__
   self.params = [p for p in self.module.parameters(self) if p.requires_grad]
File "/home/l/B/src/cv.py", line 22, in <listcomp>
   self.params = [p for p in self.module.parameters(self) if p.requires_grad]
File "/home/l/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 739, in parameters
    for name, param in self.named_parameters():
AttributeError: 'Train' object has no attribute 'named_parameters'