Block if using pytorch-summary

[navi2000]

I have been having some blocking issues. I isolated the problem in this 30 lines code:

import numpy as np
import torch
from torchsummary import summary
import torch.multiprocessing as mp
from torch.nn import functional as F

class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.l1 = torch.nn.Linear(64, 100)
        self.last = torch.nn.Linear(100, 4)
    def forward(self,x):
        y = F.relu(self.l1(x))
        result = F.log_softmax(self.last(y),dim=0)
        return result


def worker(t, worker_model):
    state = torch.from_numpy(np.random.rand(1, 64)).float()
    value = worker_model(state)
    print ("Never arrive here")


model = MyModel()
model.share_memory()
#summary(model, (1,64))

p = mp.Process(target=worker, args=(0, model))
p.start() 
p.join()

If you uncomment the 'summary' line, the worker function will be blocked in the call to worker_model(state).