densechen / penal-connection

The source code of penal connection, which was first proposed in `Rethinking skip connection model as a learnable Markov chain`.

Hi, 👋 I'm Dengsheng Chen

penal-connection's Issues

TypeError: unsupported operand type(s) for +: 'Tensor' and 'RemovableHandle'

Dear authors,
It's an interesting work that improves residual networks with the Markov chain.
However when I give the Network an input 'torch.randn(16,1)' and run the test.py file, it reports the following error.

Traceback (most recent call last):
  File "/code/penal-connection/test.py", line 31, in <module>
    print(Net()(torch.randn(16,1)))
  File "/root/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/code/penal-connection/test.py", line 23, in forward
    fc_2 = fc_1 + self.fc_2(fc_1)
TypeError: unsupported operand type(s) for +: 'Tensor' and 'RemovableHandle'
(base) root@hd7e87ca34f84105ac629d3bd34ac54f-task0-0:/code/penal-connection# python test.py
Traceback (most recent call last):
  File "/code/penal-connection/test.py", line 31, in <module>
    print(Net()(torch.randn(16,1)))
  File "/root/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/code/penal-connection/test.py", line 23, in forward
    fc_2 = fc_1 + self.fc_2(fc_1)
TypeError: unsupported operand type(s) for +: 'Tensor' and 'RemovableHandle'

And here is my test.py file:

import penal_connection as pc
import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc_1 = nn.Linear(1, 4)
        
        # use as a layer
        self.fc_2 = nn.Sequential(
            nn.Linear(4, 4),
            pc.PenalConnection(tau=1e-4),
        )

        self.fc_3 = nn.Linear(4, 4)

        self.fc_4 = nn.Linear(4, 1)
    
    def forward(self, x):
        fc_1 = self.fc_1(x)

        fc_2 = fc_1 + self.fc_2(fc_1)

        # use as a function
        fc_3 = fc_2 + pc.penal_connection(self.fc_3(fc_2), tau=1e-4)

        return self.fc_4(fc_3)

print(Net()(torch.randn(16,1)))