I tried to create 3d CNN using Pytorch. the following code works with 5 images but does not work with 336 images.
The error is: RuntimeError: $ Torch: not enough memory: you tried to allocate 166GB. Buy new RAM! at /opt/conda/conda-bld/pytorch_1550813258230/work/aten/src/TH/THGeneral.cpp:201
Can anyone help me please  ?
    def __init__(self):
        super(CNNModel, self).__init__()
        
        self.conv_layer1 = self._conv_layer_set(3, 32) 
                                                      
        self.conv_layer2 = self._conv_layer_set(32, 64)
        self.fc1 = nn.Linear(64*28*28*28, 2)
        self.fc2 = nn.Linear(1404928, num_classes)
        self.relu = nn.LeakyReLU()
        self.batch=nn.BatchNorm1d(2)
        self.drop=nn.Dropout(p=0.15, inplace = True)   
        
    def _conv_layer_set(self, in_c, out_c):
        conv_layer = nn.Sequential(
        nn.Conv3d(in_c, out_c, kernel_size=(3, 3, 3), padding=0),
        nn.LeakyReLU(),
        nn.MaxPool3d((2, 2, 2)),
        )
        return conv_layer
    

    def forward(self, x):
        # Set 1
        out = self.conv_layer1(x)
        out = self.conv_layer2(out)
        out = out.view(out.size(0), -1)
        out = self.fc1(out)
        out = self.relu(out)
        out = self.batch(out)
        out = self.drop(out)
        out = F.softmax(out, dim=1)
        return out

#Definition of hyperparameters
n_iters = 2
num_epochs = 2
# Create CNN
model = CNNModel()
#model.cuda()
print(model)
# Cross Entropy Loss 
error = nn.CrossEntropyLoss()
# SGD Optimizer
learning_rate = 0.001
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

...

#train
#Definition of hyperparameters
n_iters = 2
num_epochs = 2
loss_list_train = []
accuracy_list_train = []

for epoch in range(num_epochs):
    #for i in range(x_train1.shape[0]):
    training_data = torch.Tensor(training_data)
    targets = torch.Tensor(targets) 
    training_data = Variable(training_data.view(336,3,120,120,120))
    labels = targets
    # Clear gradients
    optimizer.zero_grad()
    # Forward propagation
    outputs = model(training_data)
    # Calculate softmax and ross entropy loss
    _, predicted = torch.max(outputs, 1)
    accuracy = accuracyCalc(predicted, targets)
    #labels = labels.tolist()
    #outputs = outputs.tolist()
    labels = labels.long() 
    labels=labels.view(-1) 
    loss = nn.CrossEntropyLoss()
    loss = loss(outputs, labels)    
    # Calculating gradients
    loss.backward()
    # Update parameters
    optimizer.step()
    loss_list_train.append(loss.data)
    accuracy_list_train.append(accuracy)
    print('Iteration: {}/{}  Loss: {}  Accuracy: {} %'.format(epoch+1,  num_epochs, loss.data, accuracy))