Unable to send pytorch model from one agent to another agent

[bangxiangyong]

from agentMET4FOF.agents import AgentMET4FOF, AgentNetwork, MonitorAgent
from agentMET4FOF.streams import SineGenerator

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


class TrainerAgent(AgentMET4FOF):
    def init_parameters(self,minus_param=0.5,plus_param=0.5):
        self.net = Net()
        self.first_time = True

    def on_received_message(self, message):
        if self.first_time:
            self.send_output(self.net)
            self.first_time = False

class SineGeneratorAgent(AgentMET4FOF):
    def init_parameters(self):
        self.stream = SineGenerator()

    def agent_loop(self):
        if self.current_state == "Running":
            sine_data = self.stream.next_sample() #dictionary
            self.send_output(sine_data['x'])

def main():
    # start agent network server
    agentNetwork = AgentNetwork()
    # init agents
    gen_agent = agentNetwork.add_agent(agentType=SineGeneratorAgent)
    train_agent = agentNetwork.add_agent(agentType=TrainerAgent)
    predict_agent = agentNetwork.add_agent(agentType=TrainerAgent)
    monitor_agent = agentNetwork.add_agent(agentType=MonitorAgent)
    # connect agents : We can connect multiple agents to any particular agent
    agentNetwork.bind_agents(gen_agent, train_agent)
    agentNetwork.bind_agents(gen_agent, predict_agent)
    # connect
    agentNetwork.bind_agents(gen_agent, monitor_agent)
    agentNetwork.bind_agents(train_agent, predict_agent)
    agentNetwork.bind_agents(predict_agent, monitor_agent)
    # set all agents states to "Running"
    agentNetwork.set_running_state()

    # allow for shutting down the network after execution
    return agentNetwork


if __name__ == '__main__':
    main()

Error message as below:

Exception in thread oneway-call: Traceback (most recent call last): File "/usr/lib/python3.7/threading.py", line 917, in _bootstrap_inner self.run() File "/home/bangxiang93/PycharmProjects/agentMET4FOF_fresh_test/venv/lib/python3.7/site-packages/Pyro4/core.py", line 1891, in run super(_OnewayCallThread, self).run() File "/usr/lib/python3.7/threading.py", line 865, in run self._target(*self._args, **self._kwargs) File "/home/bangxiang93/PycharmProjects/agentMET4FOF_fresh_test/venv/lib/python3.7/site-packages/osbrain/agent.py", line 1680, in run self._loop() File "/home/bangxiang93/PycharmProjects/agentMET4FOF_fresh_test/venv/lib/python3.7/site-packages/osbrain/agent.py", line 1208, in _loop if self._iterate(): File "/home/bangxiang93/PycharmProjects/agentMET4FOF_fresh_test/venv/lib/python3.7/site-packages/osbrain/agent.py", line 1245, in _iterate self._process_events(events) File "/home/bangxiang93/PycharmProjects/agentMET4FOF_fresh_test/venv/lib/python3.7/site-packages/osbrain/agent.py", line 1282, in _process_events self._process_single_event(socket) File "/home/bangxiang93/PycharmProjects/agentMET4FOF_fresh_test/venv/lib/python3.7/site-packages/osbrain/agent.py", line 1296, in _process_single_event self._process_sub_event(socket, address, data) File "/home/bangxiang93/PycharmProjects/agentMET4FOF_fresh_test/venv/lib/python3.7/site-packages/osbrain/agent.py", line 1452, in _process_sub_event handler(self, message) File "/home/bangxiang93/PycharmProjects/agentMET4FOF_fresh_test/agentMET4FOF-develop/agentMET4FOF/agents.py", line 334, in handle_process_data proc_msg = self.on_received_message(message) File "/home/bangxiang93/PycharmProjects/agentMET4FOF_fresh_test/agentMET4FOF-develop/agentMET4FOF_tutorials/tutorial_pytorch.py", line 35, in on_received_message self.send_output(self.net) File "/home/bangxiang93/PycharmProjects/agentMET4FOF_fresh_test/agentMET4FOF-develop/agentMET4FOF/agents.py", line 263, in send_output self.send(self.PubAddr, packed_data, topic='data') File "/home/bangxiang93/PycharmProjects/agentMET4FOF_fresh_test/venv/lib/python3.7/site-packages/osbrain/agent.py", line 1493, in send topic=topic) File "/home/bangxiang93/PycharmProjects/agentMET4FOF_fresh_test/venv/lib/python3.7/site-packages/osbrain/agent.py", line 1501, in _send_address serializer=address.serializer) File "/home/bangxiang93/PycharmProjects/agentMET4FOF_fresh_test/venv/lib/python3.7/site-packages/osbrain/agent.py", line 85, in serialize_message return _serialize_callsserializer File "/home/bangxiang93/PycharmProjects/agentMET4FOF_fresh_test/venv/lib/python3.7/site-packages/osbrain/agent.py", line 50, in 'pickle': lambda message: pickle.dumps(message, -1), _pickle.PicklingError: Can't pickle <class 'main.Net'>: it's not the same object as main.Net