def test_drop(t):
q = minitorch.dropout(t, 0.0)
idx = q._tensor.sample()
assert q[idx] == t[idx]
q = minitorch.dropout(t, 1.0)
assert q[q._tensor.sample()] == 0.0
q = minitorch.dropout(t, 1.0, ignore=True)
idx = q._tensor.sample()
assert q[idx] == t[idx]
def forward(self, x):
self.mid = self.conv1.forward(x).relu()
# print(self.mid.shape)
self.out = self.conv2.forward(self.mid).relu()
# print(self.out.shape)
pool = minitorch.avgpool2d(self.out, (4, 4))
# print(pool.shape)
pool = pool.view(BATCH, 392)
# print(pool.shape)
h = self.layer1.forward(pool).relu()
# print(h.shape)
if self.mode == "train":
h = minitorch.dropout(h, 0.25)
return minitorch.logsoftmax(self.layer2.forward(h), dim=1)
