Update gnn_layers.py

3 years ago · df75883615
parent 33c6d8e9a9
commit df75883615
1 changed files with 1 additions and 1 deletions
--- a/timm/models/layers/gnn_layers.py
+++ b/timm/models/layers/gnn_layers.py
@ -4,9 +4,9 @@ import numpy as np
 import torch
 from torch import nn
 import torch.nn.functional as F
+from timm.models.fx_features import register_notrace_module
 from .drop import DropPath
 from .pos_embed import build_sincos2d_pos_embed
-from .fx_features import register_notrace_module


 def pairwise_distance(x, y):