Update README.md

README.md

@@ -21,13 +21,12 @@ from huggingface_hub import snapshot_download
 from os.path import join as pjoin
 from safetensors import safe_open
 
-#Simple FP4 dequant matmul
 @torch.compile(fullgraph=True)
 def matmul_fp4(x, W_q, scales, group_size, fp4_values):
     def unpack_over_cols(W_q_packed, W_nbits, num_output_cols, dtype):
         n_rows, n_cols = W_q_packed.shape
-        device, dtype  = W_q_packed.device, W_q_packed.dtype
-        shifts         = torch.arange(num_output_cols // n_cols, device=device, dtype=dtype) * W_nbits 
+        device         = W_q_packed.device
+        shifts         = torch.arange(num_output_cols // n_cols, device=device, dtype=W_q_packed.dtype) * W_nbits 
         W_q_unpacked   = ((W_q_packed.unsqueeze(-1) >> shifts) & ((1 << W_nbits) - 1)).to(dtype)
         W_q_unpacked   = W_q_unpacked.view(n_rows, num_output_cols)
         return W_q_unpacked
@@ -38,6 +37,7 @@ def matmul_fp4(x, W_q, scales, group_size, fp4_values):
     return torch.matmul(x, W_r)
 
 class AutoModelForCausalLMFP4:
+
     @classmethod
     def from_pretrained(
         cls,
@@ -48,6 +48,7 @@ class AutoModelForCausalLMFP4:
         *args,
         **kwargs
     ):
+
         #Download snapshot
         if os.path.exists(save_dir_or_hub):
             save_dir = save_dir_or_hub