Commit 4bbd4184 authored by Andrei-Claudiu Roibu's avatar Andrei-Claudiu Roibu 🖥
Browse files

increased X & y dimensions to match 3D convolution req input

parent d3b668b0
......@@ -167,6 +167,10 @@ class Solver():
X = sampled_batch[0].type(torch.FloatTensor)
y = sampled_batch[1].type(torch.LongTensor)
# We add an extra dimension (~ number of channels) for the 3D convolutions.
X = torch.unsqueeze(X, dim= 1)
y = torch.unsqueeze(y, dim= 1)
if model.is_cuda():
X = X.cuda(self.device, non_blocking=True)
y = y.cuda(self.device, non_blocking=True)
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment