Hello,
I have defined a densenet architecture in PyTorch to use it on training data consisting of 15000 samples of 128x128 images. Here is the code:
class Dense_Block(nn.Module):
    def __init__(self, in_channels):
        super(Dense_Block, self).__init__()

        self.relu = nn.ReLU(inplace = True)
        self.bn = nn.BatchNorm2d(num_features = in_channels)

        self.conv1 = nn.Conv2d(in_channels = in_channels, out_channels = 32, kernel_size = 3, stride = 1, padding = 1)
        self.conv2 = nn.Conv2d(in_channels = 32, out_channels = 32, kernel_size = 3, stride = 1, padding = 1)
        self.conv3 = nn.Conv2d(in_channels = 64, out_channels = 32, kernel_size = 3, stride = 1, padding = 1)
        self.conv4 = nn.Conv2d(in_channels = 96, out_channels = 32, kernel_size = 3, stride = 1, padding = 1)
        self.conv5 = nn.Conv2d(in_channels = 128, out_channels = 32, kernel_size = 3, stride = 1, padding = 1)

    def forward(self, x):

        bn = self.bn(x)
        conv1 = self.relu(self.conv1(bn))

        conv2 = self.relu(self.conv2(conv1))
        c2_dense = self.relu(torch.cat([conv1, conv2], 1))

        conv3 = self.relu(self.conv3(c2_dense))
        c3_dense = self.relu(torch.cat([conv1, conv2, conv3], 1))

        conv4 = self.relu(self.conv4(c3_dense))
        c4_dense = self.relu(torch.cat([conv1, conv2, conv3, conv4], 1))

        conv5 = self.relu(self.conv5(c4_dense))
        c5_dense = self.relu(torch.cat([conv1, conv2, conv3, conv4, conv5], 1))

        return c5_dense

class Transition_Layer(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Transition_Layer, self).__init__()

        self.relu = nn.ReLU(inplace = True)
        self.bn = nn.BatchNorm2d(num_features = out_channels)
        self.conv = nn.Conv2d(in_channels = in_channels, out_channels = out_channels, kernel_size = 1, bias = False)
        self.avg_pool = nn.AvgPool2d(kernel_size = 2, stride = 2, padding = 0)

    def forward(self, x):

        bn = self.bn(self.relu(self.conv(x)))
        out = self.avg_pool(bn)

        return out

class DenseNet(nn.Module):
    def __init__(self, nr_classes):
        super(DenseNet, self).__init__()

        self.lowconv = nn.Conv2d(in_channels = 1, out_channels = 64, kernel_size = 7, padding = 3, bias = False)
        self.relu = nn.ReLU()

        # Make Dense Blocks
        self.denseblock1 = self._make_dense_block(Dense_Block, 64)
        self.denseblock2 = self._make_dense_block(Dense_Block, 128)
        self.denseblock3 = self._make_dense_block(Dense_Block, 128)

        # Make transition Layers
        self.transitionLayer1 = self._make_transition_layer(Transition_Layer, in_channels = 160, out_channels = 128)
        self.transitionLayer2 = self._make_transition_layer(Transition_Layer, in_channels = 160, out_channels = 128)
        self.transitionLayer3 = self._make_transition_layer(Transition_Layer, in_channels = 160, out_channels = 64)

        # Classifier
        self.bn = nn.BatchNorm2d(num_features = 64)
        self.pre_classifier = nn.Linear(64*16*16, 512)
        self.classifier = nn.Linear(512, nr_classes)

    def _make_dense_block(self, block, in_channels):
        layers = []
        layers.append(block(in_channels))
        return nn.Sequential(*layers)

    def _make_transition_layer(self, layer, in_channels, out_channels):
        modules = []
        modules.append(layer(in_channels, out_channels))
        return nn.Sequential(*modules)

    def forward(self, x):
        out = self.relu(self.lowconv(x))

        out = self.denseblock1(out)
        out = self.transitionLayer1(out)

        out = self.denseblock2(out)
        out = self.transitionLayer2(out)

        out = self.denseblock3(out)
        out = self.transitionLayer3(out)

        out = self.bn(out)
#         print(out.shape)
        out = out.reshape(-1, 64*16*16)

        out = self.pre_classifier(out)
        out = self.classifier(out)

        return out

Then I define my DataSet class:
class MyDataset(Dataset):
    def __init__(self, images, n, labels=None, transforms=None):
        self.X = images
        self.y = labels
        self.n = n
        self.transforms = transforms
         
    def __len__(self):
        return (len(self.X))
    
    def __getitem__(self, i):
        data = self.X.iloc[i, :]
        data = np.asarray(data).astype(np.float).reshape(1,n,n)
        
        if self.transforms:
            data = self.transforms(data).reshape(1,n,n)
            
        if self.y is not None:
            y = self.y.iloc[i,:]
            y = np.asarray(y).astype(np.float).reshape(128,) # for 128-vector of labels
            return (data, y)
        else:
            return data

Then I create the instances of the train, dev, and test data:
train_data = MyDataset(train_images, n, train_labels, None)
dev_data = MyDataset(dev_images, n, dev_labels, None)
test_data = MyDataset(test_images, n, test_labels, None)

The shapes of  train_images ,  dev_images  and  test_images  are respectively  (15000, 16384) ,  (4000, 16384)  and  (1000, 16384) . So there are in total 20000 samples of 128x128 (=16384) images.
The shapes of  train_labels ,  dev_labels  and  test_labels  are respectively  (15000, 128) ,  (4000, 128)  and  (1000, 128) . So there are in total 20000 samples of 128 vectors.
I define also a custom loss function:
# Huber
def Huber(yHat,y,delta=1.):
    n_samples = yHat.size()[0]
    n_points = yHat.size()[1]
    preds = yHat
    labels = y
    size = yHat.size()[0]*yHat.size()[1]
    diff = yHat - y
    return torch.sum(torch.where(torch.abs(diff) < delta,.5*diff**2 , delta*(torch.abs(diff)-.5*delta**2))) / size

Then I create an instance of the model:
densenet = DenseNet(nr_classes=128).float().to('cuda:0')

Then I initialize parameters, create train- and dev-set dataloaders, and train the model using Adam optimizer and Huber loss-function:
def main():
    nn.init.kaiming_uniform_(list(densenet.parameters())[0], nonlinearity = 'relu')
    loader = DataLoader(train_data,batch_size=200,shuffle=False,num_workers=0)
    loader_dev = DataLoader(dev_data,batch_size=10,shuffle=None,num_workers=0)
    N_epochs = 10
    for epoch in range(N_epochs):
          optimizer = optim.Adam(densenet.parameters(), lr=.001, betas=(0.9, 0.999), eps=1e-08)
          for batch in loader:
                images = batch[0].float().to('cuda:0')
                labels = batch[1].float().to('cuda:0')
                preds = densenet(images)
                loss = Loss(preds,labels).Huber()
            
                with torch.no_grad():
                    loss_dev = 0
                    for batch_dev in loader_dev:
                        images_dev = batch_dev[0].float().to('cuda:0')
                        labels_dev = batch_dev[1].float().to('cuda:0')
                        preds_dev = densenet(images_dev)
                        loss_ = Loss(preds_dev,labels_dev).Huber()
                        loss_dev += loss_

                 optimizer.zero_grad()
                 loss.backward()
                 optimizer.step()

if __name__ == '__main__':
    multiprocessing.freeze_support()
    main()

I have two GPUs identified as cuda:0 (~ 24 Gb memory) and cuda:1 (~ 6 Gb memory)
With cuda:1 device I get the error message:

  File "D:\Jupiter_playground\fashion_mnist_tidied.py", line 1127, in <module>
    main()

  File "D:\Jupiter_playground\fashion_mnist_tidied.py", line 1074, in main
    preds = network(images) # Pass Batch

  File "C:\Users\Admin\.conda\envs\pytorch_env\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)

  File "D:\Jupiter_playground\fashion_mnist_tidied.py", line 752, in forward
    out = self.denseblock1(out)

  File "C:\Users\Admin\.conda\envs\pytorch_env\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)

  File "C:\Users\Admin\.conda\envs\pytorch_env\lib\site-packages\torch\nn\modules\container.py", line 117, in forward
    input = module(input)

  File "C:\Users\Admin\.conda\envs\pytorch_env\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)

  File "D:\Jupiter_playground\fashion_mnist_tidied.py", line 691, in forward
    c3_dense = self.relu(torch.cat([conv1, conv2, conv3], 1))

RuntimeError: CUDA out of memory. Tried to allocate 1.17 GiB (GPU 1; 6.00 GiB total capacity; 4.34 GiB already allocated; 16.62 MiB free; 4.34 GiB reserved in total by PyTorch)

Then I tried to switch to cuda:0 device which has much more memory capacity, and the error in this case reads as:

  File "D:\Jupiter_playground\fashion_mnist_tidied.py", line 1127, in <module>
    main()

  File "D:\Jupiter_playground\fashion_mnist_tidied.py", line 1074, in main
    preds = network(images) # Pass Batch

  File "C:\Users\Admin\.conda\envs\pytorch_env\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)

  File "D:\Jupiter_playground\fashion_mnist_tidied.py", line 752, in forward
    out = self.denseblock1(out)

  File "C:\Users\Admin\.conda\envs\pytorch_env\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)

  File "C:\Users\Admin\.conda\envs\pytorch_env\lib\site-packages\torch\nn\modules\container.py", line 117, in forward
    input = module(input)

  File "C:\Users\Admin\.conda\envs\pytorch_env\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)

  File "D:\Jupiter_playground\fashion_mnist_tidied.py", line 691, in forward
    c3_dense = self.relu(torch.cat([conv1, conv2, conv3], 1))

RuntimeError: CUDA out of memory. Tried to allocate 1.17 GiB (GPU 0; 24.00 GiB total capacity; 21.59 GiB already allocated; 372.94 MiB free; 21.69 GiB reserved in total by PyTorch)

Why does PyTorch allocate almost all available memory?
However, when I use train-set of 6 images and dev-set of 3 images (test-set of 1 image), training with cuda-devices works fine.