Distill: add feature map

cifar10-fast-simple/model.py

@@ -132,10 +132,11 @@ class ResNetBagOfTricks(nn.Module):
         x = self.conv6(x)
         x = self.conv7(x)
         x = x + self.conv9(self.conv8(x))
+        feature_map = x
         x = self.pool10(x)
         x = x.reshape(x.size(0), x.size(1))
         x = self.linear11(x)
         x = self.scale_out * x
-        return x
+        return x, feature_map
 
 Model = ResNetBagOfTricks

cifar10-fast-simple/studentmodel.py

@@ -1,3 +1,4 @@
+import torch
 import torch.nn as nn