# Unable to access custom model functions when wrapping it with torch.prepare\_model
**URL:** https://discuss.ray.io/t/unable-to-access-custom-model-functions-when-wrapping-it-with-torch-prepare-model/9716
**Category:** Uncategorized
**Created:** [March 10, 2023, 10:28am UTC](https://discuss.ray.io/t/unable-to-access-custom-model-functions-when-wrapping-it-with-torch-prepare-model/9716 "2023-03-10T10:28:22Z")
**Posts on this page:** 5
**Page:** 1
### Author:  [@AxelN](https://discuss.ray.io/u/AxelN)
#### Post date: [March 10, 2023, 10:28am UTC](https://discuss.ray.io/t/unable-to-access-custom-model-functions-when-wrapping-it-with-torch-prepare-model/9716/1 "2023-03-10T10:28:23Z")
We use:
ray 2.3.0
torch: 1.13.1
We have a model that have custom functions, example:
```python
class ExampleNetwork(torch.nn.Module):
def forward(self):
...
def custom_function_1(self):
...
def custom_function_2(self):
...
model = ray.train.torch.prepare_model(ExampleNetwork())
```
When you use the function ray.train.torch.prepare\_model, the resulting model works differently depending on the number of workers we use in the TorchTrainer (ray.train.torch.TorchTrainer).
If we only use one worker, the model do not need to be parallelized and is therefore still a nn.Module, but if we use multiple workers it needs to be parallised and instead become a DistributedDataParallel (torch.nn.parallel.DistributedDataParallel).
The problem with this is that when the model is parallelized, the prepare\_model-function do not wrap the custom functions, and to access these you need to change the calls from:
model.custom\_function → model.module.custom\_function
The standard functions like forward still works as intended.
It is not really scalable to add checks for the type of the model before every function call, so I am wondering if there is another way to prepare the model or a way to wrap the resulting model so we do not need to have different training code dependent on the number of workers we are going to be training on.
---
### Author:  [@gjoliver](https://discuss.ray.io/u/gjoliver)
#### Post date: [March 13, 2023, 7:57pm UTC](https://discuss.ray.io/t/unable-to-access-custom-model-functions-when-wrapping-it-with-torch-prepare-model/9716/2 "2023-03-13T19:57:19Z")
Hi, this is a really good question.
As you said, we by-pass the entire wrapping logic if world\_size \<= 1:
>
I will discuss this with the team internally and see how we can make this part of the experience better.
thanks again for the feedback.
---
### Author:  [@AxelN](https://discuss.ray.io/u/AxelN)
#### Post date: [March 14, 2023, 7:53am UTC](https://discuss.ray.io/t/unable-to-access-custom-model-functions-when-wrapping-it-with-torch-prepare-model/9716/3 "2023-03-14T07:53:12Z")
Thanks for he response and appreciate you taking this further.
One temporary solution I have found so far, but have not tested out extensively so might bring unintended consequences, is to do:
```python
wrapped_model = ray.train.torch.prepare_model(ExampleNetwork())
if ray.air.session.get_world_size() > 1:
model = wrapped_model.module
else:
model = wrapped_model
```
---
### Author:  [@gjoliver](https://discuss.ray.io/u/gjoliver)
#### Post date: [March 14, 2023, 8:08am UTC](https://discuss.ray.io/t/unable-to-access-custom-model-functions-when-wrapping-it-with-torch-prepare-model/9716/4 "2023-03-14T08:08:12Z")
yeah, seems like a nice workaround
---
### Author:  [@Jules\_Damji](https://discuss.ray.io/u/Jules_Damji)
#### Post date: [March 14, 2023, 4:54pm UTC](https://discuss.ray.io/t/unable-to-access-custom-model-functions-when-wrapping-it-with-torch-prepare-model/9716/5 "2023-03-14T16:54:56Z")
Thanks @gjoliver. @AxelN , it seems you all sorted.