diff --git a/CMakeLists.txt b/CMakeLists.txt
index 511059d..3adf434 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -13,6 +13,7 @@ set(LIBRARY_SOURCES
     src/layers/dense.cpp
 )
 
+set(CMAKE_CUDA_ARCHITECTURES 75)
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} -arch=sm_75)
 
diff --git a/include/layers/dense.h b/include/layers/dense.h
index fd89707..93073d4 100644
--- a/include/layers/dense.h
+++ b/include/layers/dense.h
@@ -27,7 +27,7 @@ namespace Layers {
         float* d_weights;
         float* d_biases;
 
-        std::vector<std::vector<float>> weights;
+        std::vector<float> weights;
         std::vector<float> biases;
 
         void initializeWeights();
diff --git a/src/layers/dense.cpp b/src/layers/dense.cpp
index c39d92a..1c8dfec 100644
--- a/src/layers/dense.cpp
+++ b/src/layers/dense.cpp
@@ -4,15 +4,13 @@
 #include <cuda_runtime.h>
 #include <cublas_v2.h>
 #include <cstdio>
-#include <stdexcept>
-	
-
+#include <random>
 
 Layers::Dense::Dense(int inputSize, int outputSize, cublasHandle_t cublasHandle)
     : inputSize(inputSize), outputSize(outputSize), cublasHandle(cublasHandle) {
 
     // Allocate memory for weights and biases
-    weights.resize(outputSize, std::vector<float>(inputSize));
+    weights.resize(outputSize * inputSize);
     biases.resize(outputSize);
 
     initializeWeights();
@@ -32,17 +30,22 @@ Layers::Dense::~Dense() {
 }
 
 void Layers::Dense::initializeWeights() {
-    for (auto& row : weights) {
-        for (float& weight : row) {
-            weight = 0.0f;
+    int numWeights = inputSize * outputSize;
+
+    std::random_device rd;
+    std::mt19937 gen(rd());
+    std::normal_distribution<float> dist(0.0f, 0.01f); // Xavier initialization
+
+    for (int i = 0; i < outputSize; ++i) {
+        for (int j = 0; j < inputSize; ++j) {
+            int idx = IDX2C(i, j, inputSize);
+            weights[idx] = dist(gen);
         }
     }
 }
 
 void Layers::Dense::initializeBiases() {
-    for (float& bias : biases) {
-        bias = 0.0f;
-    }
+    std::fill(biases.begin(), biases.end(), 0.1f);
 }
 
 void Layers::Dense::forward(const float* d_input, float* d_output) {
@@ -58,12 +61,20 @@ void Layers::Dense::toCuda() {
     CUBLAS_CHECK(cublasSetVector(biases.size(), sizeof(float), biases.data(), 1, d_biases, 1));
 }
 
-void Layers::Dense::setWeights(const std::vector<std::vector<float>>& weights) {
-    this->weights = weights;
+void Layers::Dense::setWeights(const std::vector<std::vector<float>>& weights_input) {
+    int numWeights = inputSize * outputSize;
+
+    for (int i = 0; i < outputSize; ++i) {
+        for (int j = 0; j < inputSize; ++j) {
+            int idx = IDX2C(i, j, inputSize);
+            weights[idx] = weights_input[i][j];
+        }
+    }
+
     toCuda();
 }
 
-void Layers::Dense::setBiases(const std::vector<float>& biases) {
-    this->biases = biases;
+void Layers::Dense::setBiases(const std::vector<float>& biases_input) {
+    std::copy(biases_input.begin(), biases_input.end(), biases.begin());
     toCuda();
 }
\ No newline at end of file