Duet Split Learning GPU devices

[umermjd11]

In split learning mnist demo

class SyNet_client(sy.Module):
    def __init__(self, torch_ref):
        super(SyNet_client, self).__init__(torch_ref=torch_ref)
        self.lin = nn.Linear(input_size, hidden_sizes[0])
    def forward(self, x):
        x = self.lin(x)
        x = self.torch_ref.nn.functional.relu(x)
        return x
class SyNet_server(sy.Module):
    def __init__(self, torch_ref):
        super(SyNet_server, self).__init__(torch_ref=torch_ref)
        self.lin2 = nn.Linear(hidden_sizes[0], hidden_sizes[1])
        self.lin3 = nn.Linear(hidden_sizes[1], output_size)
        self.sft = nn.LogSoftmax(dim=1)     
    def forward(self, x):
        x = self.lin2(x)
        x = nn.ReLU()(x)
        x = self.lin3(x)
        x = self.sft(x)
        return x

#Model 1
model1 = SyNet_client(torch)
model1_ptr = model1.send(duet)
remote_device = duet.torch.device("cuda" if duet.torch.cuda.is_available() else "cpu")
#Model 2
model2 = SyNet_server(torch)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

How can we send model1 and model2 to their GPU devices?
I tried with following code but it gives error

model1_ptr.to(remote_device)
model2.to(device)

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-25-e534ab5b9982> in <module>
----> 1 model1_ptr.to(remote_device)
      2 model2.to(device)

/opt/conda/lib/python3.8/site-packages/syft/lib/torch/module.py in __getattr__(self, name)
    141                 return modules[name]
    142 
--> 143         return object.__getattribute__(self, name)
    144 
    145     def train(self, mode: bool = True) -> "Module":

AttributeError: 'SyNet_client' object has no attribute 'to'

[znreza]

Hi @umermjd11 You can send your model to use "cuda" (after checking if it has "cuda") by doing so

if has_cuda:
    model.cuda(device)
else:
    model.cpu()

Check this example for details