'Proxy' object does not support item assignment

username: KeyKy
cooked: <pre><code class="lang-auto">class TemporalShift(nn.Module):
    def __init__(self, n_segment=3, n_div=8):
        super(TemporalShift, self).__init__()
        self.n_segment = n_segment
        self.fold_div = n_div

    def forward(self, x):
        x = self.shift(x, self.n_segment, fold_div=self.fold_div)
        return x

    @staticmethod
    def shift(x, n_segment, fold_div=3):
        nt, c, h, w = x.size()
        n_batch = nt // n_segment
        x = x.view(n_batch, n_segment, c, h, w)

        fold = c // fold_div

        #out = torch.zeros_like(x)
        out = x.clone().zero_()
        out[:, :-1, :fold] = x[:, 1:, :fold]  # shift left
        out[:, 1:, fold: 2 * fold] = x[:, :-1, fold: 2 * fold]  # shift right
        out[:, :, 2 * fold:] = x[:, :, 2 * fold:]  # not shift

        return out.view(nt, c, h, w)
</code></pre>
<p>when using fx to tracing tsm, i get this error:</p>
<p>out[:, :-1, :fold] = x[:, 1:, :fold]<br>
TypeError: ‘Proxy’ object does not support item assignment</p>

username: Vasiliy_Kuznetsov
cooked: <p>Hi <a class="mention" href="/u/keyky">@keyky</a>, this is a limitation of symbolic tracing with FX.  Here is a workaround using torch.fx.wrap:</p>
<pre><code class="lang-auto">@torch.fx.wrap                                                           
def shift_left(out, x, fold):                                            
    out[:, :-1, :fold] = x[:, 1:, :fold]  # shift left                   
                                                                         
@torch.fx.wrap                                                           
def shift_right(out, x, fold):                                           
    out[:, 1:, fold: 2 * fold] = x[:, :-1, fold: 2 * fold]  # shift right
                                                                         
@torch.fx.wrap                                                           
def not_shift(out, x, fold):                                             
    out[:, :, 2 * fold:] = x[:, :, 2 * fold:]  # not shift               
                                                                         
class TemporalShift(nn.Module):                                          
    def __init__(self, n_segment=3, n_div=8):                            
        super(TemporalShift, self).__init__()                            
        self.n_segment = n_segment                                       
        self.fold_div = n_div                                            
                                                                         
    def forward(self, x):                                                
        x = self.shift(x, self.n_segment, fold_div=self.fold_div)        
        return x                                                         
                                                                         
    @staticmethod                                                        
    def shift(x, n_segment, fold_div=3):                                 
        nt, c, h, w = x.size()                                           
        n_batch = nt // n_segment                                        
        x = x.view(n_batch, n_segment, c, h, w)                          
                                                                         
        fold = c // fold_div                                             
                                                                         
        #out = torch.zeros_like(x)                                       
        out = x.clone().zero_()                                          
        shift_left(out, x, fold)                                         
        shift_right(out, x, fold)                                        
        not_shift(out, x, fold)                                          
                                                                         
        return out.view(nt, c, h, w)                                     
                                                                         
m = TemporalShift()                                                      
ms = torch.fx.symbolic_trace(m)                                          
</code></pre>