From bca9d0b8e506e8fc5c31126ab873ef97c8d44ac1 Mon Sep 17 00:00:00 2001
From: Akemi Izuko <akemi@noway.moe>
Date: Sat, 23 Nov 2024 23:42:14 -0700
Subject: [PATCH] Distill: add feature map

---
 cifar10-fast-simple/model.py        | 3 ++-
 cifar10-fast-simple/studentmodel.py | 1 +
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/cifar10-fast-simple/model.py b/cifar10-fast-simple/model.py
index fe33495..602849c 100644
--- a/cifar10-fast-simple/model.py
+++ b/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
diff --git a/cifar10-fast-simple/studentmodel.py b/cifar10-fast-simple/studentmodel.py
index 071ffff..076f2cb 100644
--- a/cifar10-fast-simple/studentmodel.py
+++ b/cifar10-fast-simple/studentmodel.py
@@ -1,3 +1,4 @@
+import torch
 import torch.nn as nn