# Using GRPOTrainer with a custom PyTorch module?
**URL:**
**Author:**  [@enyoukai](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fenyoukai)%5C
**Post date:** [April 21, 2025, 4:30am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fusing-grpotrainer-with-a-custom-pytorch-module%2F151526%2F1 "2025-04-21T04:30:11Z")
Hello! I was wondering what changes I would need to make to use GRPOTrainer with a custom PyTorch module class.
Currently I have an nn.Module subclass that wraps around an existing Huggingface transformer, except with a custom forward and generate function.
I was wondering though if there were resources on either converting an nn.Module to a transformer to be used with Trainer, or what other functionality I would need to implement as well as changes I’d need to make to my forward and generate methods to work with
---
**Author:**  [@John6666](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FJohn6666)%5C
**Post date:** [April 21, 2025, 6:46am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fusing-grpotrainer-with-a-custom-pytorch-module%2F151526%2F2 "2025-04-21T06:46:10Z")
If you inherit **PreTrainedModel** , you should have most of the necessary functions. As for Trainer, it seems that you can modify loss functions, gradient-related functions, and so on.
> [@What is difference between Model.register\_for\_auto\_class and AutoModel.register?](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fwhat-is-difference-between-model-register-for-auto-class-and-automodel-register%2F29866)%3A
>
> I have created some new model and want to register it, so I can use AutoModelFor…() to call it. I have read [Sharing custom models](https://huggingface.co/docs/transformers/custom_models), but I can’t tell the difference between ResnetConfig.register\_for\_auto\_class() ResnetModel.register\_for\_auto\_class("AutoModel") ResnetModelForImageClassification.register\_for\_auto\_class("AutoModelForImageClassification") and AutoConfig.register("resnet", ResnetConfig) AutoModel.register(ResnetConfig, ResnetModel) AutoModelForImageClassification.register(ResnetCon…
> **[Adding a new model to Transformers](https://huggingface.co/docs/transformers/add_new_model)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.
---
**Author:**  [@enyoukai](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fenyoukai)%5C
**Post date:** [April 29, 2025, 4:33pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fusing-grpotrainer-with-a-custom-pytorch-module%2F151526%2F3 "2025-04-29T16:33:55Z")
Sorry for bumping but do you know if Trainer supports passing in inputs\_embeds for generate yet? I ended up just monkey-patching the generate method to use my own generation code, but it seems that when I directly pass in inputs\_embeds instead of input\_ids into the original generate method, I get the following error:
```auto
Traceback (most recent call last):
File "test_grpo.py", line 124, in
trainer.train()
File "env/lib64/python3.9/site-packages/transformers/trainer.py", line 2241, in train
return inner_training_loop(
File "env/lib64/python3.9/site-packages/transformers/trainer.py", line 2548, in _inner_training_loop
tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
File "env/lib64/python3.9/site-packages/transformers/trainer.py", line 3692, in training_step
inputs = self._prepare_inputs(inputs)
File "env/lib64/python3.9/site-packages/trl/trainer/grpo_trainer.py", line 576, in _prepare_inputs
eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)]
IndexError: argmax(): Expected reduction dim 1 to have non-zero size.
Traceback (most recent call last):
File "test_grpo.py", line 124, in
trainer.train()
File "env/lib64/python3.9/site-packages/transformers/trainer.py", line 2241, in train
return inner_training_loop(
File "env/lib64/python3.9/site-packages/transformers/trainer.py", line 2548, in _inner_training_loop
tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
File "env/lib64/python3.9/site-packages/transformers/trainer.py", line 3692, in training_step
inputs = self._prepare_inputs(inputs)
File "env/lib64/python3.9/site-packages/trl/trainer/grpo_trainer.py", line 576, in _prepare_inputs
eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)]
IndexError: argmax(): Expected reduction dim 1 to have non-zero size.
```
---
**Author:**  [@enyoukai](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fenyoukai)%5C
**Post date:** [April 29, 2025, 4:44pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fusing-grpotrainer-with-a-custom-pytorch-module%2F151526%2F4 "2025-04-29T16:44:39Z")
fixed. turns out i just need to pad back the output with the original prompt lol